Rewrite determinant of diagonal matrix as product of diagonal

pytensor/tensor/rewriting/linalg.py

@@ -5,12 +5,14 @@
 from pytensor import Variable
 from pytensor.graph import Apply, FunctionGraph
 from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
-from pytensor.tensor.basic import TensorVariable, diagonal
+from pytensor.scalar.basic import Mul
+from pytensor.tensor.basic import Eye, TensorVariable, diagonal
 from pytensor.tensor.blas import Dot22
 from pytensor.tensor.blockwise import Blockwise
-from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.math import Dot, Prod, _matrix_matrix_matmul, log, prod
 from pytensor.tensor.nlinalg import (
+    Det,
     KroneckerProduct,
     MatrixInverse,
     MatrixPinv,
@@ -377,3 +379,55 @@ def local_lift_through_linalg(
                 return [block_diag(*inner_matrices)]
             else:
                 raise NotImplementedError  # pragma: no cover
+
+
+# Det Diag Rewrite
+@register_canonicalize
+@register_stabilize
+@node_rewriter([det])
+def det_diag_rewrite(fgraph, node):
+    # Find if we have Blockwise Op
+    if not (isinstance(node.op, Blockwise) and isinstance(node.op.core_op, Det)):
+        return None
+
+    node_input_op = node.inputs[0]
+    # Check if the op is Elemwise and mul
+    if not (
+        node_input_op.owner is not None
+        and isinstance(node_input_op.owner.op, Elemwise)
+        and isinstance(node_input_op.owner.op.scalar_op, Mul)
+    ):
+        return None
+
+    # Find whether any of the inputs to mul is Eye
+    inputs_to_mul = node_input_op.owner.inputs
+    eye_check = False
+    eye_input = []
+    for mul_input in inputs_to_mul:
+        if isinstance(mul_input.owner.op, Eye):
+            eye_check = True
+            eye_input.append(mul_input)
+            break  # This is so that we only find the first one? Can be changed
+    if not (eye_check):
+        return None
+
+    # Get all non Eye inputs (scalars/matrices/vectors)
+    non_eye_inputs = list(set(inputs_to_mul) - set(eye_input))
+
+    # Dealing with only one other input
+    if len(non_eye_inputs) >= 2:
+        raise NotImplementedError
+
+    # Checking if original x was matrix
+    if not non_eye_inputs[0].owner:
+        # It was a matrix
+        det_val = non_eye_inputs[0].diagonal().prod()
+    else:
+        # Check for scalar (ndim = 0) or vector (ndim = 1)
+        sca_vec_input = non_eye_inputs[0].owner.inputs[0]
+        if sca_vec_input.ndim == 0:
+            det_val = sca_vec_input ** eye_input[0].shape[0]
+        else:
+            det_val = sca_vec_input.prod()
+
+    return [det_val]

tests/tensor/rewriting/test_linalg.py

@@ -390,3 +390,37 @@ def test_local_lift_through_linalg(constructor, f_op, f, g_op, g):
     test_vals = [x @ np.swapaxes(x, -1, -2) for x in test_vals]
 
     np.testing.assert_allclose(f1(*test_vals), f2(*test_vals), atol=1e-8)
+
+
+@pytest.mark.parametrize(
+    "shape", [(), (7,), (7, 7)], ids=["scalar", "vector", "matrix"]
+)
+def test_det_diag_rewrite(shape):
+    # Initializing x based on scalar/vector/matrix
+    x = pt.tensor("x", shape=shape)
+    y = pt.eye(7) * x
+    # Calculating determinant value using pt.linalg.det
+    z_det = pt.linalg.det(y)
+
+    # REWRITE TEST
+    f_rewritten = function([x], z_det, mode="FAST_RUN")
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    assert not any(isinstance(node.op, Det) for node in nodes)
+
+    # NUMERIC VALUE TEST
+    f_det = function([x], z_det)
+    if len(shape) == 0:
+        x_test = np.random.rand()
+    elif len(shape) == 1:
+        x_test = np.random.rand(7)
+    else:
+        x_test = np.random.rand(7, 7)
+    det_val = f_det(x_test)
+    rewritten_val = f_rewritten(x_test)
+
+    assert_allclose(
+        det_val,
+        rewritten_val,
+        atol=1e-4 if config.floatX == "float32" else 1e-8,
+        rtol=1e-4 if config.floatX == "float32" else 1e-8,
+    )