Sylvester rewrite

Author: jessegrabowski

pytensor/tensor/_linalg/solve/rewriting.py

@@ -8,18 +8,23 @@
 from pytensor.scan.rewriting import scan_seqopt1
 from pytensor.tensor._linalg.solve.linear_control import (
     SolveBilinearDiscreteLyapunov,
+    SolveSylvester,
     solve_discrete_lyapunov,
 )
 from pytensor.tensor._linalg.solve.tridiagonal import (
     tridiagonal_lu_factor,
     tridiagonal_lu_solve,
 )
-from pytensor.tensor.basic import atleast_Nd
+from pytensor.tensor.basic import atleast_Nd, diagonal
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import DimShuffle
-from pytensor.tensor.rewriting.basic import register_specialize
+from pytensor.tensor.rewriting.basic import (
+    register_canonicalize,
+    register_specialize,
+    register_stabilize,
+)
 from pytensor.tensor.rewriting.blockwise import blockwise_of
-from pytensor.tensor.rewriting.linalg import is_matrix_transpose
+from pytensor.tensor.rewriting.linalg import _is_diagonal, is_matrix_transpose
 from pytensor.tensor.slinalg import Solve, cho_solve, cholesky, lu_factor, lu_solve
 from pytensor.tensor.variable import TensorVariable
 
@@ -306,3 +311,27 @@ def jax_bilinear_lyapunov_to_direct(fgraph, node):
     "jax",
     position=0.9,  # Run before canonicalization
 )
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([blockwise_of(SolveSylvester)])
+def rewrite_sylvester_diag_to_elemwise(fgraph, node):
+    """solve_sylvester(diag_A, diag_B, C) -> C / (a[:, None] + b[None, :])
+
+    When both coefficient matrices are diagonal, :math:`AX + XB = C` decouples into
+    :math:`X_{ij} = C_{ij} / (a_i + b_j)`.
+    """
+    A, B, C = node.inputs
+
+    if not _is_diagonal(A, fgraph) or not _is_diagonal(B, fgraph):
+        return None
+
+    diag_a = diagonal(A, axis1=-2, axis2=-1)
+    diag_b = diagonal(B, axis1=-2, axis2=-1)
+    denom = diag_a[..., :, None] + diag_b[..., None, :]
+
+    X_out = node.outputs[0]
+    new_x = (C / denom).astype(X_out.type.dtype)
+
+    return [new_x]