added all tests

Author: tanish1729

pytensor/tensor/nlinalg.py

@@ -266,7 +266,6 @@ def __str__(self):
         return "SLogDet"
 
 
-# slogdet = Blockwise(SLogDet())
 def slogdet(x):
     det_val = det(x)
     return ptm.sign(det_val), ptm.log(ptm.abs(det_val))

pytensor/tensor/rewriting/linalg.py

@@ -969,50 +969,34 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
     return [eye_input * (non_eye_input**0.5)]
 
 
-def _check_log_abs_det(fgraph, client):
-    # First, we find abs
-    if not (isinstance(client.op, Elemwise) and isinstance(client.op.scalar_op, Abs)):
-        return False
-
-    # Check whether log is a client of abs
-    for client_2 in fgraph.clients[client.outputs[0]]:
-        if not (
-            isinstance(client_2[0].op, Elemwise)
-            and isinstance(client_2[0].op.scalar_op, Log)
-        ):
-            return False
-
-    return True
-
-
 @register_specialize
 @node_rewriter([det])
 def slogdet_specialization(fgraph, node):
     replacements = {}
-    for client in fgraph.clients[node.outputs[0]]:
+    for client, _ in fgraph.clients[node.outputs[0]]:
         # Check for sign(det)
-        if isinstance(client[0].op, Elemwise) and isinstance(
-            client[0].op.scalar_op, Sign
-        ):
+        if isinstance(client.op, Elemwise) and isinstance(client.op.scalar_op, Sign):
             x = node.inputs[0]
             sign_det_x, slog_det_x = SLogDet()(x)
-            replacements[client[0].outputs[0]] = sign_det_x
+            replacements[client.outputs[0]] = sign_det_x
 
         # Check for log(abs(det))
-        elif _check_log_abs_det(fgraph, client[0]):
-            x = node.inputs[0]
-            sign_det_x, slog_det_x = SLogDet()(x)
-            replacements[fgraph.clients[client[0].outputs[0]][0][0].outputs[0]] = (
-                slog_det_x
-            )
+        elif isinstance(client.op, Elemwise) and isinstance(client.op.scalar_op, Abs):
+            for client_2, _ in fgraph.clients[client.outputs[0]]:
+                if isinstance(client_2.op, Elemwise) and isinstance(
+                    client_2.op.scalar_op, Log
+                ):
+                    x = node.inputs[0]
+                    sign_det_x, slog_det_x = SLogDet()(x)
+                    replacements[fgraph.clients[client.outputs[0]][0][0].outputs[0]] = (
+                        slog_det_x
+                    )
 
         # Check for log(det)
-        elif isinstance(client[0].op, Elemwise) and isinstance(
-            client[0].op.scalar_op, Log
-        ):
+        elif isinstance(client.op, Elemwise) and isinstance(client.op.scalar_op, Log):
             x = node.inputs[0]
             sign_det_x, slog_det_x = SLogDet()(x)
-            replacements[client[0].outputs[0]] = pt.where(
+            replacements[client.outputs[0]] = pt.where(
                 pt.eq(sign_det_x, -1), np.nan, slog_det_x
             )
 

tests/tensor/rewriting/test_linalg.py

@@ -781,7 +781,7 @@ def test_det_kronecker_rewrite():
     a, b = pt.dmatrices("a", "b")
     kron_prod = pt.linalg.kron(a, b)
     det_output = pt.linalg.det(kron_prod)
-    f_rewritten = function([kron_prod], [det_output], mode="FAST_RUN")
+    f_rewritten = function([a, b], [det_output], mode="FAST_RUN")
 
     # Rewrite Test
     nodes = f_rewritten.maker.fgraph.apply_nodes
@@ -791,7 +791,7 @@ def test_det_kronecker_rewrite():
     a_test, b_test = np.random.rand(2, 20, 20)
     kron_prod_test = np.kron(a_test, b_test)
     det_output_test = np.linalg.det(kron_prod_test)
-    rewritten_det_val = f_rewritten(kron_prod_test)
+    rewritten_det_val = f_rewritten(a_test, b_test)
     assert_allclose(
         det_output_test,
         rewritten_det_val,
@@ -800,6 +800,35 @@ def test_det_kronecker_rewrite():
     )
 
 
+def test_slogdet_kronecker_rewrite():
+    a, b = pt.dmatrices("a", "b")
+    kron_prod = pt.linalg.kron(a, b)
+    sign_output, logdet_output = pt.linalg.slogdet(kron_prod)
+    f_rewritten = function([a, b], [sign_output, logdet_output], mode="FAST_RUN")
+
+    # Rewrite Test
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    assert not any(isinstance(node.op, KroneckerProduct) for node in nodes)
+
+    # Value Test
+    a_test, b_test = np.random.rand(2, 20, 20)
+    kron_prod_test = np.kron(a_test, b_test)
+    sign_output_test, logdet_output_test = np.linalg.slogdet(kron_prod_test)
+    rewritten_sign_val, rewritten_logdet_val = f_rewritten(a_test, b_test)
+    assert_allclose(
+        sign_output_test,
+        rewritten_sign_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
+    assert_allclose(
+        logdet_output_test,
+        rewritten_logdet_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
+
+
 def test_cholesky_eye_rewrite():
     x = pt.eye(10)
     L = pt.linalg.cholesky(x)
@@ -904,20 +933,60 @@ def test_rewrite_cholesky_diag_to_sqrt_diag_not_applied():
 
 
 def test_slogdet_specialisation():
-    x = pt.dmatrix("x")
-    det_x = pt.linalg.det(x)
-    log_abs_det_x = pt.log(pt.abs(det_x))
-    sign_det_x = pt.sign(det_x)
+    x, a = pt.dmatrix("x"), np.random.rand(20, 20)
+    det_x, det_a = pt.linalg.det(x), np.linalg.det(a)
+    log_abs_det_x, log_abs_det_a = pt.log(pt.abs(det_x)), np.log(np.abs(det_a))
+    log_det_x, log_det_a = pt.log(det_x), np.log(det_a)
+    sign_det_x, sign_det_a = pt.sign(det_x), np.sign(det_a)
     exp_det_x = pt.exp(det_x)
+    # REWRITE TESTS
     # sign(det(x))
     f = function([x], [sign_det_x], mode="FAST_RUN")
     nodes = f.maker.fgraph.apply_nodes
-    assert any(isinstance(node.op, SLogDet) for node in nodes)
+    assert len([node for node in nodes if isinstance(node.op, SLogDet)]) == 1
+    assert not any(isinstance(node.op, Det) for node in nodes)
+    rw_sign_det_a = f(a)
+    assert_allclose(
+        sign_det_a,
+        rw_sign_det_a,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
     # log(abs(det(x)))
     f = function([x], [log_abs_det_x], mode="FAST_RUN")
     nodes = f.maker.fgraph.apply_nodes
-    assert any(isinstance(node.op, SLogDet) for node in nodes)
+    assert len([node for node in nodes if isinstance(node.op, SLogDet)]) == 1
+    assert not any(isinstance(node.op, Det) for node in nodes)
+    rw_log_abs_det_a = f(a)
+    assert_allclose(
+        log_abs_det_a,
+        rw_log_abs_det_a,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
+    # log(det(x))
+    f = function([x], [log_det_x], mode="FAST_RUN")
+    nodes = f.maker.fgraph.apply_nodes
+    assert len([node for node in nodes if isinstance(node.op, SLogDet)]) == 1
+    assert not any(isinstance(node.op, Det) for node in nodes)
+    rw_log_det_a = f(a)
+    assert_allclose(
+        log_det_a,
+        rw_log_det_a,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )
+    # more than 1 valid function
+    f = function([x], [sign_det_x, log_abs_det_x], mode="FAST_RUN")
+    nodes = f.maker.fgraph.apply_nodes
+    assert len([node for node in nodes if isinstance(node.op, SLogDet)]) == 1
+    assert not any(isinstance(node.op, Det) for node in nodes)
     # other functions (rewrite shouldnt be applied to these)
+    # only invalid functions
     f = function([x], [exp_det_x], mode="FAST_RUN")
     nodes = f.maker.fgraph.apply_nodes
     assert not any(isinstance(node.op, SLogDet) for node in nodes)
+    # invalid + valid function
+    f = function([x], [exp_det_x, sign_det_x], mode="FAST_RUN")
+    nodes = f.maker.fgraph.apply_nodes
+    assert not any(isinstance(node.op, SLogDet) for node in nodes)