added rewrite for diag(block_diag)

Author: tanish1729

pytensor/tensor/rewriting/linalg.py

@@ -12,8 +12,11 @@
 from pytensor.scalar.basic import Mul
 from pytensor.tensor.basic import (
     AllocDiag,
+    ExtractDiag,
     Eye,
     TensorVariable,
+    concatenate,
+    diag,
     diagonal,
 )
 from pytensor.tensor.blas import Dot22
@@ -701,3 +704,24 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
         non_eye_input = pt.shape_padaxis(non_eye_diag, -2)
 
     return [eye_input / non_eye_input]
+
+
+@register_canonicalize
+@register_stabilize
+@node_rewriter([ExtractDiag])
+def rewrite_diag_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
+    submatrices = potential_blockdiag.inputs
+    submatrices_diag = [diag(submatrices[i]) for i in range(len(submatrices))]
+    output = [concatenate(submatrices_diag)]
+
+    return output

tests/tensor/rewriting/test_linalg.py

@@ -662,3 +662,30 @@ def test_inv_diag_from_diag(inv_op):
         atol=ATOL,
         rtol=RTOL,
     )
+
+
+def test_diag_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)])
+    diag_output = pt.diag(bd_output)
+    f_rewritten = function([sub_matrices], diag_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)]
+    )
+    diag_output_test = np.diag(bd_output_test)
+    rewritten_val = f_rewritten(sub_matrices_test)
+    assert_allclose(
+        diag_output_test,
+        rewritten_val,
+        atol=1e-3 if config.floatX == "float32" else 1e-8,
+        rtol=1e-3 if config.floatX == "float32" else 1e-8,
+    )