Infer dtype from node.outputs.type.dtype

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -43,11 +43,11 @@
     register_stabilize,
 )
 from pytensor.tensor.slinalg import (
+    BilinearSolveDiscreteLyapunov,
     BlockDiagonal,
     Cholesky,
     Solve,
     SolveBase,
-    _solve_bilinear_discrete_lyapunov,
     block_diag,
     cholesky,
     solve,
@@ -972,11 +972,14 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
     return [eye_input * (non_eye_input**0.5)]
 
 
-@node_rewriter([_solve_bilinear_discrete_lyapunov])  #  type: ignore
+@node_rewriter([Blockwise])  #  type: ignore
 def jax_bilinaer_lyapunov_to_direct(fgraph: FunctionGraph, node: Apply):
     """
     Replace BilinearSolveDiscreteLyapunov with a direct computation that is supported by JAX
     """
+    if not isinstance(node.op.core_op, BilinearSolveDiscreteLyapunov):
+        return None
+
     A, B = (cast(TensorVariable, x) for x in node.inputs)
     result = solve_discrete_lyapunov(A, B, method="direct")
 

pytensor/tensor/slinalg.py

@@ -797,7 +797,8 @@ def make_node(self, A, B):
     def perform(self, node, inputs, output_storage):
         (A, B) = inputs
         X = output_storage[0]
-        out_dtype = pytensor.scalar.upcast(A.dtype, B.dtype)
+        out_dtype = node.outputs[0].type.dtype
+
         X[0] = scipy.linalg.solve_continuous_lyapunov(A, B).astype(out_dtype)
 
     def infer_shape(self, fgraph, node, shapes):
@@ -843,9 +844,9 @@ def perform(self, node, inputs, output_storage):
         (A, B) = inputs
         X = output_storage[0]
 
-        dtype = pytensor.scalar.upcast(A.dtype, B.dtype)
+        out_dtype = node.outputs[0].type.dtype
         X[0] = scipy.linalg.solve_discrete_lyapunov(A, B, method="bilinear").astype(
-            dtype
+            out_dtype
         )
 
     def infer_shape(self, fgraph, node, shapes):
@@ -868,12 +869,6 @@ def grad(self, inputs, output_grads):
         return [A_bar, Q_bar]
 
 
-_solve_continuous_lyapunov = Blockwise(SolveContinuousLyapunov())
-_solve_bilinear_discrete_lyapunov = cast(
-    typing.Callable, Blockwise(BilinearSolveDiscreteLyapunov())
-)
-
-
 def _direct_solve_discrete_lyapunov(
     A: TensorVariable, Q: TensorVariable
 ) -> TensorVariable:
@@ -937,7 +932,7 @@ def solve_discrete_lyapunov(
         return cast(TensorVariable, X)
 
     elif method == "bilinear":
-        return cast(TensorVariable, _solve_bilinear_discrete_lyapunov(A, Q))
+        return cast(TensorVariable, Blockwise(BilinearSolveDiscreteLyapunov())(A, Q))
 
     else:
         raise ValueError(f"Unknown method {method}")
@@ -961,10 +956,10 @@ def solve_continuous_lyapunov(A: TensorLike, Q: TensorLike) -> TensorVariable:
 
     """
 
-    return cast(TensorVariable, _solve_continuous_lyapunov(A, Q))
+    return cast(TensorVariable, Blockwise(SolveContinuousLyapunov())(A, Q))
 
 
-class SolveDiscreteARE(pt.Op):
+class SolveDiscreteARE(Op):
     __props__ = ("enforce_Q_symmetric",)
     gufunc_signature = "(m,m),(m,n),(m,m),(n,n)->(m,m)"
 

tests/link/jax/test_slinalg.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Literal
 
 import numpy as np
@@ -208,11 +209,13 @@ def test_jax_solve_discrete_lyapunov(
     out = pt_slinalg.solve_discrete_lyapunov(A, B, method=method)
     out_fg = FunctionGraph([A, B], [out])
 
+    atol = rtol = 1e-8 if config.floatX == "float64" else 1e-3
     compare_jax_and_py(
         out_fg,
         [
             np.random.normal(size=shape).astype(config.floatX),
             np.random.normal(size=shape).astype(config.floatX),
         ],
         jax_mode="JAX",
+        assert_fn=partial(np.testing.assert_allclose, atol=atol, rtol=rtol),
     )

tests/tensor/test_slinalg.py

@@ -537,8 +537,8 @@ def test_solve_discrete_lyapunov(
         precision = int(dtype[-2:])  # 64 or 32
         dtype = f"complex{int(2 * precision)}"
 
-    A1, A2 = rng.normal(size=(2, *shape)).astype(dtype)
-    Q1, Q2 = rng.normal(size=(2, *shape)).astype(dtype)
+    A1, A2 = rng.normal(size=(2, *shape))
+    Q1, Q2 = rng.normal(size=(2, *shape))
 
     if use_complex:
         A = A1 + 1j * A2
@@ -547,6 +547,8 @@ def test_solve_discrete_lyapunov(
         A = A1
         Q = Q1
 
+    A, Q = A.astype(dtype), Q.astype(dtype)
+
     a = pt.tensor(name="a", shape=shape, dtype=dtype)
     q = pt.tensor(name="q", shape=shape, dtype=dtype)
 
@@ -585,15 +587,20 @@ def test_solve_discrete_lyapunov_gradient(
 @pytest.mark.parametrize("shape", [(5, 5), (5, 5, 5)], ids=["matrix", "batched"])
 @pytest.mark.parametrize("use_complex", [False, True], ids=["float", "complex"])
 def test_solve_continuous_lyapunov(shape: tuple[int], use_complex: bool):
+    dtype = config.floatX
+    if use_complex and dtype == "float32":
+        pytest.skip(
+            "Not enough precision in complex64 to do schur decomposition "
+            "(ill-conditioned matrix errors arise)"
+        )
     rng = np.random.default_rng(utt.fetch_seed())
 
-    dtype = config.floatX
     if use_complex:
         precision = int(dtype[-2:])  # 64 or 32
         dtype = f"complex{int(2 * precision)}"
 
-    A1, A2 = rng.normal(size=(2, *shape)).astype(dtype)
-    Q1, Q2 = rng.normal(size=(2, *shape)).astype(dtype)
+    A1, A2 = rng.normal(size=(2, *shape))
+    Q1, Q2 = rng.normal(size=(2, *shape))
 
     if use_complex:
         A = A1 + 1j * A2
@@ -602,9 +609,13 @@ def test_solve_continuous_lyapunov(shape: tuple[int], use_complex: bool):
         A = A1
         Q = Q1
 
+    A, Q = A.astype(dtype), Q.astype(dtype)
+
     a = pt.tensor(name="a", shape=shape, dtype=dtype)
     q = pt.tensor(name="q", shape=shape, dtype=dtype)
-    f = function([a, q], [solve_continuous_lyapunov(a, q)])
+    x = solve_continuous_lyapunov(a, q)
+
+    f = function([a, q], x)
 
     X = f(A, Q)