Improved Linear Control Ops, with Numba Dispatches #1840

jessegrabowski · 2026-01-12T05:14:13Z

Description

Another day, another PR with a huge pile of matrix decompositions nobody asked for :)

solve_discrete_lyapunov is quite important for statespace, because it's used to initialize a (big) covariance matrix. There are two modes for solve_discrete_lyapunov: direct and bilinear. We have bilinear support in the C backend, but Numba and JAX only support direct. Direct involves making a kronecker product of a (big) covariance matrix, so it sucks.

Several linear control Ops: solve_sylvester, solve_discrete_lyapunov, and solve_continuous_lyapunov, are all actually the same Op. They involve a Schur decomposition of the inputs, followed by a call to the TRSYL routine. Aside from that, it's just some dots and solves.

This PR adds the Schur and TRSYL Ops, then uses these two components to implement all of the linear control ops (except for solve_discrete_are, that's another kettle of fish). These two new Ops don't have gradients, but the solvers that use them do, and lop_override is used to handle this. TRSYL is also not added to __all__, as it is not intended to be user-facing. Schur is a well-known matrix decomposition, so I exposed that one.

This has several advantages:

By exposing the inner graphs (that were previously in scipy.linalg) we can do linag rewrites
Less code to maintain -- all of the linear control Ops now flow through solve_sylvester, so I removed all of the specialized Ops/gradients for the others
Free backend dispatching -- as long as we can implement Schur and TRSYL in a backend, we get all the rest of the stuff "for free".

As it turns out, we can implement Schur in Numba and JAX, and TRSYL in Numba. This PR also does that.

I think @Dekermanjian added some code to JAX that would let us dispatch these ops to jax as well, but I don't know if it's merged/released. I would be willing to do that in this PR if it's out.

Related Issue

Closes #
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

Copilot

Pull request overview

This PR refactors and improves linear control operations in PyTensor by introducing Schur and TRSYL as fundamental building blocks. The refactoring consolidates solve_sylvester, solve_continuous_lyapunov, and solve_discrete_lyapunov to flow through a common implementation using these new Ops, enabling better graph optimization and native backend support for Numba and JAX.

Changes:

Introduces Schur decomposition Op and TRSYL (Sylvester equation solver) as new building blocks
Refactors solve_sylvester, solve_continuous_lyapunov, and solve_discrete_lyapunov to use the new Ops
Adds Numba and JAX dispatches for the new Ops with comprehensive test coverage

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
pytensor/tensor/slinalg.py	Core implementation of TRSYL, Schur, and refactored linear control solvers with gradient support
pytensor/tensor/rewriting/linalg.py	Updated rewriting rules to reference renamed `SolveBilinearDiscreteLyapunov` Op
pytensor/link/numba/dispatch/slinalg.py	Numba dispatches for Schur and TRSYL with proper dtype handling
pytensor/link/numba/dispatch/linalg/_LAPACK.py	Low-level LAPACK bindings for gees (Schur) and trsyl
pytensor/link/jax/dispatch/slinalg.py	JAX dispatch for Schur decomposition
tests/tensor/test_slinalg.py	Comprehensive tests for solve_sylvester, refactored solve_continuous_lyapunov tests, and new Schur tests
tests/link/numba/test_slinalg.py	Numba-specific tests for Schur with overwrite and sort parameter validation
tests/link/jax/test_slinalg.py	JAX-specific tests for Schur

pytensor/tensor/slinalg.py

ricardoV94

This looks pretty good, comments are pretty minor, besides what the bot found

ricardoV94 · 2026-01-12T10:47:19Z

pytensor/link/numba/dispatch/linalg/decomposition/schur.py

+        VS = np.empty((_LDVS, _N), dtype=dtype)
+        RWORK = np.empty(_N, dtype=w_type)
+        BWORK = val_to_int_ptr(1)
+        INFO = val_to_int_ptr(1)


Info isn't used?

ricardoV94 · 2026-01-12T10:47:35Z

pytensor/link/jax/dispatch/slinalg.py

+    if op.sort is not None:
+        raise NotImplementedError("jax.scipy.linalg.schur only supports sort=None.")


Instead of failing, issue a warning that it is ignored?

ricardoV94 · 2026-01-12T10:48:05Z

pytensor/link/numba/dispatch/linalg/solve/linear_control.py

+            INFO,
+        )
+
+        if int_ptr_to_val(INFO) < 0:


> 0 never happens or it's fine?

It can be 1, which means there were multiple eigenvalues of the same value, so a solution was found by jittering the matrix. A solution is returned in that case, so we can pass it back. For reference, scipy spams warnings in this case but also gives you back a matrix.

ricardoV94 · 2026-01-12T10:48:34Z

pytensor/link/numba/dispatch/linalg/solve/linear_control.py

+            return np.full_like(C_copy, np.nan)
+
+        # CC now contains the solution, scale it
+        X = SCALE * C_copy


Reuse the array and then return it?

Suggested change

X = SCALE * C_copy

C_copy *= SCALE

ricardoV94 · 2026-01-12T10:50:12Z

pytensor/link/numba/dispatch/slinalg.py

+    sort = op.sort
+
+    if sort is not None:
+        raise NotImplementedError(


Same, warn, or fallback to obj mode, don't fail

pytensor/tensor/slinalg.py

ricardoV94 · 2026-01-12T10:54:37Z

pytensor/tensor/slinalg.py

+        AxA = vec_kron(A, A.conj())
+        eye = pt.eye(AxA.shape[-1])
+
+        vec_Q = Q.reshape((*Q.shape[:-2], -1))


join_dims cof cof

ricardoV94 · 2026-01-12T10:55:21Z

pytensor/tensor/slinalg.py

+        sort: Literal["lhp", "rhp", "iuc", "ouc"] | None = None,
+    ):
+        self.output = output
+        self.gufunc_signature = "(m,m)->(m,m),(m,m)"


if gufunc_signature is constant, define it as an Op attribute, instead of in __init__

ricardoV94 · 2026-01-12T10:58:36Z

tests/tensor/test_slinalg.py

+    ids=["float", "complex", "batch_float"],
+)
+def test_solve_continuous_sylvester(shape: tuple[int], use_complex: bool):
+    # batch-complex case got an error from BatchedDot not implemented for complex numbers


then the rewrite that introduces BatchedDot shouldn't try to introduce it?

ricardoV94 · 2026-01-12T10:59:51Z

tests/tensor/test_slinalg.py

+        precision = int(dtype[-2:])  # 64 or 32
+        dtype = f"complex{int(2 * precision)}"


Suggested change

precision = int(dtype[-2:]) # 64 or 32

dtype = f"complex{int(2 * precision)}"

dtype = "complex128" if dtype == "float64" else "complex64"

Dekermanjian · 2026-01-12T12:59:16Z

I think @Dekermanjian added some code to JAX that would let us dispatch these ops to jax as well, but I don't know if it's merged/released. I would be willing to do that in this PR if it's out.

Hey @jessegrabowski, yes I did. I added the Sylvester solver. I believe you can implement the Lyapunov like so:

import jax
import jax.numpy as jnp
import jax.scipy as jsp

def solve_discrete_lyapunov(A, C, method="bilinear"):
    if method == "bilinear":
        eye = jnp.eye(A.shape[0])
        aH = A.conj().transpose()
        aHI_inv = jnp.linalg.inv(aH + eye)
        b = jnp.dot(aH - eye, aHI_inv)
        c = 2 * jnp.dot(jnp.dot(jnp.linalg.inv(A + eye), C), aHI_inv)
        solution = jsp.linalg.solve_sylvester(b.conj().transpose(), b.conj(), -c)
        return solution
    if method == "direct":
        lhs = jnp.kron(A, A.conj())
        lhs = jnp.eye(lhs.shape[0]) - lhs
        x = jsp.linalg.solve(lhs, C.flatten())
        return jnp.reshape(x, C.shape)

I would also be happy to open a PR and add this method directly to Jax.

jessegrabowski · 2026-01-12T16:20:01Z

I would also be happy to open a PR and add this method directly to Jax.

If we have jsp.linalg.solve_sylvester I will just dispatch that, then we don't need a lyapunov dispatch at all. This PR removes the lyapunov Op, and essentially just does that code you posted as a pytensor graph (except using solve instead of inv), so we would be able to dispatch that whole graph, potentially with optimizations once we can infer matrix structures.

Dekermanjian · 2026-01-12T16:33:46Z

I would also be happy to open a PR and add this method directly to Jax.

If we have jsp.linalg.solve_sylvester I will just dispatch that, then we don't need a lyapunov dispatch at all. This PR removes the lyapunov Op, and essentially just does that code you posted as a pytensor graph (except using solve instead of inv), so we would be able to dispatch that whole graph, potentially with optimizations once we can infer matrix structures.

This is really super cool! Please let me know what I can do to help out!!

jessegrabowski · 2026-01-12T16:35:52Z

Drag that monster statespace PR over the finish line :P

ricardoV94 · 2026-01-13T02:16:45Z

pytensor/tensor/slinalg.py



-class SolveContinuousLyapunov(Op):
+class TRSYL(Op):


There's a tensor._linalg module, don't be shy to start populating it instead of these nlinalg / slinalg

jessegrabowski added enhancement New feature or request jax numba SciPy compatibility linalg Linear algebra labels Jan 12, 2026

jessegrabowski requested review from Copilot and ricardoV94 January 12, 2026 05:16

Copilot started reviewing on behalf of jessegrabowski January 12, 2026 05:16 View session

jessegrabowski force-pushed the numba-control branch from 969a7e0 to 32e7fd3 Compare January 12, 2026 05:19

Copilot AI reviewed Jan 12, 2026

View reviewed changes

ricardoV94 reviewed Jan 12, 2026

View reviewed changes

jessegrabowski added 6 commits January 12, 2026 19:07

Add Schur decomposition Op

1c960fc

Add Numba dispatch for Schur

16798ff

Add JAX dispatch for Schur

4831af7

Add TRSYL Op, refactor linear control Ops

e198c2a

Numba dispatch for linear control Ops

e946541

JAX dispatch for linear control Ops

1040325

jessegrabowski force-pushed the numba-control branch from 32e7fd3 to 1040325 Compare January 13, 2026 01:57

ricardoV94 reviewed Jan 13, 2026

View reviewed changes

Move linear control Ops to linear_control.py

abca0f7

		if op.sort is not None:
		raise NotImplementedError("jax.scipy.linalg.schur only supports sort=None.")

		precision = int(dtype[-2:]) # 64 or 32
		dtype = f"complex{int(2 * precision)}"

	precision = int(dtype[-2:]) # 64 or 32
	dtype = f"complex{int(2 * precision)}"
	dtype = "complex128" if dtype == "float64" else "complex64"



		class SolveContinuousLyapunov(Op):
		class TRSYL(Op):

Improved Linear Control Ops, with Numba Dispatches #1840

Are you sure you want to change the base?

Improved Linear Control Ops, with Numba Dispatches #1840

Uh oh!

Conversation

jessegrabowski commented Jan 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Dekermanjian commented Jan 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jessegrabowski commented Jan 12, 2026

Uh oh!

Dekermanjian commented Jan 12, 2026

Uh oh!

jessegrabowski commented Jan 12, 2026

Uh oh!

ricardoV94 Jan 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jessegrabowski commented Jan 12, 2026 •

edited

Loading

Dekermanjian commented Jan 12, 2026 •

edited

Loading

ricardoV94 Jan 13, 2026 •

edited

Loading