paramatrize cholesky_ldotlt test

Author: dehorsley

pytensor/sandbox/linalg/ops.py

@@ -112,8 +112,11 @@ def psd_solve_with_chol(fgraph, node):
 @register_canonicalize
 @register_stabilize
 @node_rewriter([Cholesky])
-def chol_of_dot_chol(fgraph, node):
+def cholesky_ldotlt(fgraph, node):
     """
+    rewrite cholesky(dot(L, L.T), lower=True) = L, where L is lower triangular,
+    or cholesky(dot(U.T, U), upper=True) = U where U is upper triangular.
+
     This utilizes a boolean `lower_triangular` or `upper_triangular` tag on matrices.
     """
     if not isinstance(node.op, Cholesky):
@@ -133,6 +136,7 @@ def chol_of_dot_chol(fgraph, node):
         and r.owner.op.new_order == (1, 0)
         and r.owner.inputs[0] == l
     ):
+        print("found right form")
         if node.op.lower:
             return [l]
         return [r]

tests/sandbox/linalg/test_linalg.py

@@ -1,11 +1,12 @@
 import numpy as np
 import numpy.linalg
+import pytest
 import scipy.linalg
 
 import pytensor
 from pytensor import function
 from pytensor import tensor as at
-from pytensor.compile import DeepCopyOp, ViewOp
+from pytensor.compile import get_default_mode
 from pytensor.configdefaults import config
 from pytensor.sandbox.linalg.ops import inv_as_solve, spectral_radius_bound
 from pytensor.tensor.elemwise import DimShuffle
@@ -156,94 +157,72 @@ def test_matrix_inverse_solve():
     assert isinstance(out.owner.op, Solve)
 
 
-def test_cholesky_dot_lower():
-    cholesky_lower = Cholesky(lower=True)
-    cholesky_upper = Cholesky(lower=False)
+@pytest.mark.parametrize("tag", ("lower", "upper", None))
+@pytest.mark.parametrize("cholesky_form", ("lower", "upper"))
+@pytest.mark.parametrize("product", ("lower", "upper", None))
+def test_cholesky_ldotlt(tag, cholesky_form, product):
+    cholesky = Cholesky(lower=(cholesky_form == "lower"))
 
-    L = matrix("L")
-    L.tag.lower_triangular = True
+    transform_removes_chol = tag is not None and product == tag
+    transform_transposes = transform_removes_chol and cholesky_form != tag
 
-    C = cholesky_lower(L.dot(L.T))
-    f = pytensor.function([L], C)
+    A = matrix("L")
+    if tag:
+        setattr(A.tag, tag + "_triangular", True)
 
-    if config.mode != "FAST_COMPILE":
-        assert (f.maker.fgraph.outputs[0] == f.maker.fgraph.inputs[0]) or (
-            (o := f.maker.fgraph.outputs[0].owner)
-            and isinstance(o.op, (DeepCopyOp, ViewOp))
-            and o.inputs[0] == f.maker.fgraph.inputs[0]
-        )
+    if product == "lower":
+        M = A.dot(A.T)
+    elif product == "upper":
+        M = A.T.dot(A)
+    else:
+        M = A
 
-    # Test some concrete value through f:
-    Lv = np.array([[2, 0], [1, 4]])
-    assert np.all(
-        np.isclose(
-            scipy.linalg.cholesky(np.dot(Lv, Lv.T), lower=True),
-            f(Lv),
-        )
-    )
+    C = cholesky(M)
+    f = pytensor.function([A], C, mode=get_default_mode().including("cholesky_ldotlt"))
 
-    # Test upper decomposition factors down to a transpose
-    C = cholesky_upper(L.dot(L.T))
-    f = pytensor.function([L], C)
-    if config.mode != "FAST_COMPILE":
-        assert (
-            (o1 := f.maker.fgraph.outputs[0].owner)
-            and isinstance(o1.op, (DeepCopyOp, ViewOp))
-            and (o2 := o1.inputs[0].owner)
-            and isinstance(o2.op, DimShuffle)
-            and o2.op.new_order == (1, 0)
-            and o2.inputs[0] == f.maker.fgraph.inputs[0]
-        )
+    print(f.maker.fgraph.apply_nodes)
 
-    assert np.all(
-        np.isclose(
-            scipy.linalg.cholesky(np.dot(Lv, Lv.T), lower=False),
-            f(Lv),
-        )
+    no_cholesky_in_graph = not any(
+        isinstance(node.op, Cholesky) for node in f.maker.fgraph.apply_nodes
     )
 
+    assert no_cholesky_in_graph == transform_removes_chol
 
-def test_cholesky_dot_upper():
-    cholesky_lower = Cholesky(lower=True)
-    cholesky_upper = Cholesky(lower=False)
-
-    U = matrix("U")
-    U.tag.upper_triangular = True
-
-    C = cholesky_upper(U.T.dot(U))
-    f = pytensor.function([U], C)
-    if config.mode != "FAST_COMPILE":
-        assert (f.maker.fgraph.outputs[0] == f.maker.fgraph.inputs[0]) or (
-            (o := f.maker.fgraph.outputs[0].owner)
-            and isinstance(o.op, (DeepCopyOp, ViewOp))
-            and o.inputs[0] == f.maker.fgraph.inputs[0]
+    if transform_transposes:
+        assert any(
+            isinstance(node.op, DimShuffle) and node.op.new_order == (1, 0)
+            for node in f.maker.fgraph.apply_nodes
         )
 
-    # Test some concrete value through f:
-    Uv = np.array([[2, 1], [0, 4]])
-    assert np.all(
-        np.isclose(
-            scipy.linalg.cholesky(np.dot(Uv.T, Uv), lower=False),
-            f(Uv),
+    # Test some concrete value through f
+    # there must be lower triangular (f assumes they are)
+    Avs = [
+        np.eye(1),
+        np.eye(10),
+        np.array([[2, 0], [1, 4]]),
+    ]
+    if not tag:
+        # these must be positive def
+        Avs.extend(
+            [
+                np.ones((4, 4)) + np.eye(4),
+            ]
         )
-    )
 
-    # Test lower decomposition factors down to a transpose
-    C = cholesky_lower(U.T.dot(U))
-    f = pytensor.function([U], C)
-    if config.mode != "FAST_COMPILE":
-        assert (
-            (o1 := f.maker.fgraph.outputs[0].owner)
-            and isinstance(o1.op, (DeepCopyOp, ViewOp))
-            and (o2 := o1.inputs[0].owner)
-            and isinstance(o2.op, DimShuffle)
-            and o2.op.new_order == (1, 0)
-            and o2.inputs[0] == f.maker.fgraph.inputs[0]
-        )
-
-    assert np.all(
-        np.isclose(
-            scipy.linalg.cholesky(np.dot(Uv.T, Uv), lower=True),
-            f(Uv),
+    for Av in Avs:
+        if tag == "upper":
+            Av = Av.T
+
+        if product == "lower":
+            Mv = Av.dot(Av.T)
+        elif product == "upper":
+            Mv = Av.T.dot(Av)
+        else:
+            Mv = Av
+
+        assert np.all(
+            np.isclose(
+                scipy.linalg.cholesky(Mv, lower=(cholesky_form == "lower")),
+                f(Av),
+            )
         )
-    )