Add jax rewrite to eliminate BilinearSolveDiscreteLyapunov

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -4,9 +4,11 @@
 
 from pytensor import Variable
 from pytensor import tensor as pt
+from pytensor.compile import optdb
 from pytensor.graph import Apply, FunctionGraph
 from pytensor.graph.rewriting.basic import (
     copy_stack_trace,
+    in2out,
     node_rewriter,
 )
 from pytensor.scalar.basic import Mul
@@ -41,10 +43,12 @@
     register_stabilize,
 )
 from pytensor.tensor.slinalg import (
+    BilinearSolveDiscreteLyapunov,
     BlockDiagonal,
     Cholesky,
     Solve,
     SolveBase,
+    _direct_solve_discrete_lyapunov,
     block_diag,
     cholesky,
     solve,
@@ -966,3 +970,30 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
             non_eye_input = pt.shape_padaxis(non_eye_input, -2)
 
     return [eye_input * (non_eye_input**0.5)]
+
+
+@node_rewriter([Blockwise])
+def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
+    """
+    Replace BilinearSolveDiscreteLyapunov with a direct computation that is supported by JAX
+    """
+
+    # Check if the op is BilinearSolveDiscreteLyapunov
+    if not isinstance(node.op.core_op, BilinearSolveDiscreteLyapunov):
+        return None
+
+    # Extract the inputs
+    (A, B) = node.inputs
+
+    # Compute the result
+    result = _direct_solve_discrete_lyapunov(A, B)
+
+    return [result]
+
+
+optdb.register(
+    "jax_bilinaer_lyapunov_to_direct",
+    in2out(jax_bilinaer_lyapunov_to_direct),
+    "jax",
+    position=100,
+)

tests/link/jax/test_slinalg.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 import numpy as np
 import pytest
 
@@ -194,3 +196,19 @@ def test_jax_eigvalsh(lower):
             None,
         ],
     )
+
+
+@pytest.mark.parametrize("method", ["direct", "bilinear"])
+def test_jax_solve_discrete_lyapunov(method: Literal["direct", "bilinear"]):
+    A = matrix("A")
+    B = matrix("B")
+    out = pt_slinalg.solve_discrete_lyapunov(A, B, method=method)
+    out_fg = FunctionGraph([A, B], [out])
+
+    compare_jax_and_py(
+        out_fg,
+        [
+            np.random.normal(size=(5, 5)).astype(config.floatX),
+            np.random.normal(size=(5, 5)).astype(config.floatX),
+        ],
+    )