added rewrite for determinant of blockdiag

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -635,3 +635,24 @@ def rewrite_diag_blockdiag(fgraph, node):
     output = [concatenate(submatrices_diag)]
 
     return output
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([det])
+def rewrite_det_blockdiag(fgraph, node):
+    # Check for inner block_diag operation
+    potential_blockdiag = node.inputs[0].owner
+    if not (
+        potential_blockdiag
+        and isinstance(potential_blockdiag.op, Blockwise)
+        and isinstance(potential_blockdiag.op.core_op, BlockDiagonal)
+    ):
+        return None
+
+    # Find the composing sub_matrices
+    sub_matrices = potential_blockdiag.inputs
+    det_sub_matrices = [det(sub_matrices[i]) for i in range(len(sub_matrices))]
+    prod_det_sub_matrices = prod(det_sub_matrices)
+
+    return [prod_det_sub_matrices]

tests/tensor/rewriting/test_linalg.py

@@ -595,3 +595,30 @@ def test_diag_blockdiag_rewrite():
         atol=1e-3 if config.floatX == "float32" else 1e-8,
         rtol=1e-3 if config.floatX == "float32" else 1e-8,
     )
+
+
+def test_det_blockdiag_rewrite():
+    n_matrices = 100
+    matrix_size = (5, 5)
+    sub_matrices = pt.tensor("sub_matrices", shape=(n_matrices, *matrix_size))
+    bd_output = pt.linalg.block_diag(*[sub_matrices[i] for i in range(n_matrices)])
+    det_output = pt.linalg.det(bd_output)
+    f_rewritten = function([sub_matrices], det_output, mode="FAST_RUN")
+
+    # Rewrite Test
+    nodes = f_rewritten.maker.fgraph.apply_nodes
+    assert not any(isinstance(node.op, BlockDiagonal) for node in nodes)
+
+    # Value Test
+    sub_matrices_test = np.random.rand(n_matrices, *matrix_size)
+    bd_output_test = scipy.linalg.block_diag(
+        *[sub_matrices_test[i] for i in range(n_matrices)]
+    )
+    det_output_test = np.linalg.det(bd_output_test)
+    rewritten_val = f_rewritten(sub_matrices_test)
+    assert_allclose(
+        det_output_test,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )