Check for SolveBase to catch all cases

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -47,8 +47,10 @@
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
+    CholeskySolve,
     Solve,
     SolveBase,
+    SolveTriangular,
     _bilinear_solve_discrete_lyapunov,
     block_diag,
     cholesky,
@@ -908,6 +910,11 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
         return None
 
     [input] = node.inputs
+
+    # Check if input is a (1, 1) matrix
+    if all(input.type.broadcastable[:-2]):
+        return [pt.sqrt(input)]
+
     # Check for use of pt.diag first
     if (
         input.owner
@@ -1031,15 +1038,28 @@ def scalar_solve_to_divison(fgraph, node):
     """
 
     core_op = node.op.core_op
-    if not isinstance(core_op, Solve):
+    if not isinstance(core_op, SolveBase):
         return None
 
     a, b = node.inputs
     old_out = node.outputs[0]
     if not all(a.broadcastable[-2:]):
         return None
 
-    new_out = b / a
+    # Special handling for different types of solve
+    match core_op:
+        case SolveTriangular():
+            # Corner case: if user asked for a triangular solve with a unit diagonal, a is taken to be 1
+            new_out = b / a if not core_op.unit_diagonal else b
+        case CholeskySolve():
+            new_out = b / a**2
+        case Solve():
+            new_out = b / a
+        case _:
+            raise NotImplementedError(
+                f"Unsupported core_op type: {type(core_op)} in scalar_solve_to_divison"
+            )
+
     if core_op.b_ndim == 1:
         new_out = new_out.squeeze(-1)
 

tests/tensor/rewriting/test_linalg.py

@@ -29,6 +29,7 @@
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
+    CholeskySolve,
     Solve,
     SolveBase,
     SolveTriangular,
@@ -995,18 +996,30 @@ def test_slogdet_specialization():
     assert not any(isinstance(node.op, SLogDet) for node in nodes)
 
 
-def test_scalar_solve_to_division_rewrite():
+@pytest.mark.parametrize(
+    "Op, fn",
+    [
+        (Solve, pt.linalg.solve),
+        (SolveTriangular, pt.linalg.solve_triangular),
+        (CholeskySolve, pt.linalg.cho_solve),
+    ],
+)
+def test_scalar_solve_to_division_rewrite(Op, fn):
     rng = np.random.default_rng(sum(map(ord, "scalar_solve_to_division_rewrite")))
 
     a = pt.dmatrix("a", shape=(1, 1))
     b = pt.dvector("b")
 
-    c = pt.linalg.solve(a, b, b_ndim=1)
+    if Op is CholeskySolve:
+        # cho_solve expects a tuple (c, lower) as the first input
+        c = fn((pt.linalg.cholesky(a), True), b, b_ndim=1)
+    else:
+        c = fn(a, b, b_ndim=1)
 
     f = function([a, b], c, mode="FAST_RUN")
     nodes = f.maker.fgraph.apply_nodes
 
-    assert not any(isinstance(node.op, Solve) for node in nodes)
+    assert not any(isinstance(node.op, Op) for node in nodes)
 
     a_val = rng.normal(size=(1, 1)).astype(pytensor.config.floatX)
     b_val = rng.normal(size=(1,)).astype(pytensor.config.floatX)