added specialised rewrite for slogdet

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -9,7 +9,7 @@
     copy_stack_trace,
     node_rewriter,
 )
-from pytensor.scalar.basic import Mul
+from pytensor.scalar.basic import Abs, Log, Mul, Sign
 from pytensor.tensor.basic import (
     AllocDiag,
     ExtractDiag,
@@ -28,6 +28,7 @@
     KroneckerProduct,
     MatrixInverse,
     MatrixPinv,
+    SLogDet,
     det,
     inv,
     kron,
@@ -964,3 +965,58 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
             non_eye_input = pt.shape_padaxis(non_eye_input, -2)
 
     return [eye_input * (non_eye_input**0.5)]
+
+
+# SLogDet Rewrites
+def check_sign_det(node):
+    if not (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Sign)):
+        return False
+
+    return True
+
+
+def check_log_abs_det(node):
+    if not (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Log)):
+        return False
+
+    potential_abs = node.inputs[0].owner
+    if not (
+        isinstance(potential_abs.op, Elemwise)
+        and isinstance(potential_abs.op.scalar_op, Abs)
+    ):
+        return False
+
+    return True
+
+
+def check_log_det(node):
+    if not (isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Log)):
+        return False
+
+    return True
+
+
+@node_rewriter(tracks=[det])
+def slogdet_specialization(fgraph, node):
+    x = node.inputs[0]
+    sign_det_x, slog_det_x = SLogDet()(x)
+    replacements = {}
+    for client in list(fgraph.clients.keys()):
+        # Check for sign(det)
+        if check_sign_det(client[0].owner):
+            replacements[client[0].owner.outputs[0]] = sign_det_x
+
+        # Check for log(abs(det))
+        elif check_log_abs_det(client[0].owner):
+            replacements[client[0].owner.outputs[0]] = slog_det_x
+
+        # Check for log(det)
+        elif check_log_det(client[0].owner):
+            pass
+            # replacements[client[0].owner.outputs[0]] = pt.where(pt.eq(sign_det_x, -1), np.nan, slog_det_x)
+
+        # Det is used directly for something else, don't rewrite to avoid computing two dets
+        else:
+            return None
+
+    return replacements or None