test concrete values and tranpose op

Author: dehorsley

tests/sandbox/linalg/test_linalg.py

@@ -1,5 +1,6 @@
 import numpy as np
 import numpy.linalg
+import scipy.linalg
 
 import pytensor
 from pytensor import function
@@ -157,21 +158,53 @@ def test_matrix_inverse_solve():
 
 def test_cholesky_dot_lower():
     cholesky_lower = Cholesky(lower=True)
+    cholesky_upper = Cholesky(lower=False)
 
     L = matrix("L")
     L.tag.lower_triangular = True
 
     C = cholesky_lower(L.dot(L.T))
     f = pytensor.function([L], 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]
         )
 
+    # 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),
+        )
+    )
+
+    # 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]
+        )
+
+    assert np.all(
+        np.isclose(
+            scipy.linalg.cholesky(np.dot(Lv, Lv.T), lower=False),
+            f(Lv),
+        )
+    )
+
 
 def test_cholesky_dot_upper():
+    cholesky_lower = Cholesky(lower=True)
     cholesky_upper = Cholesky(lower=False)
 
     U = matrix("U")
@@ -185,3 +218,32 @@ def test_cholesky_dot_upper():
             and isinstance(o.op, (DeepCopyOp, ViewOp))
             and o.inputs[0] == f.maker.fgraph.inputs[0]
         )
+
+    # 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 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),
+        )
+    )