Apply JAX rewrite before canonicalization

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -43,15 +43,15 @@
     register_stabilize,
 )
 from pytensor.tensor.slinalg import (
-    BilinearSolveDiscreteLyapunov,
     BlockDiagonal,
     Cholesky,
     Solve,
     SolveBase,
-    _direct_solve_discrete_lyapunov,
+    _solve_bilinear_discrete_lyapunov,
     block_diag,
     cholesky,
     solve,
+    solve_discrete_lyapunov,
     solve_triangular,
 )
 
@@ -972,21 +972,13 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
     return [eye_input * (non_eye_input**0.5)]
 
 
-@node_rewriter([Blockwise])
+@node_rewriter([_solve_bilinear_discrete_lyapunov])
 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 = (cast(TensorVariable, x) for x in node.inputs)
-
-    # Compute the result
-    result = _direct_solve_discrete_lyapunov(A, B)
+    result = solve_discrete_lyapunov(A, B, method="direct")
 
     return [result]
 
@@ -995,5 +987,5 @@ def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
     "jax_bilinaer_lyapunov_to_direct",
     in2out(jax_bilinaer_lyapunov_to_direct),
     "jax",
-    position=100,
+    position=0.9,  # Run before canonicalization
 )