Add remove_useless_block_diag rewrite

Author: jessegrabowski

pytensor/tensor/rewriting/linalg.py

@@ -821,6 +821,30 @@ def rewrite_inv_diag_to_diag_reciprocal(fgraph, node):
     return [eye_input / non_eye_input]
 
 
+@register_canonicalize
+@register_stabilize
+@node_rewriter([blockwise_of(BlockDiagonal)])
+def remove_useless_blockdiag(fgraph, node):
+    """
+     This rewrite takes advantage of the fact that block_diag with only one input is the input itself.
+
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        Function graph being optimized
+    node: Apply
+        Node of the function graph to be optimized
+
+    Returns
+    -------
+    list of Variable, optional
+        List of optimized variables, or None if no optimization was performed
+    """
+    if len(node.inputs) != 1:
+        return None
+    return [node.inputs[0]]
+
+
 @register_canonicalize
 @register_stabilize
 @node_rewriter([ExtractDiag])

tests/tensor/rewriting/test_linalg.py

@@ -786,6 +786,28 @@ def test_det_blockdiag_rewrite():
     )
 
 
+def test_remove_useless_blockdiag_rewrite():
+    A = pt.matrix("A", shape=(2, 2))
+
+    # Use the raw Op to dodge the eager return in pt.linalg.block_diag when n_inputs=1
+    bd = Blockwise(BlockDiagonal(n_inputs=1))(A)
+    bd.dprint()
+    fg = FunctionGraph(inputs=[A], outputs=[bd])
+    assert (
+        any(
+            isinstance(node.op, BlockDiagonal)
+            or (
+                isinstance(node.op, Blockwise)
+                and isinstance(node.op.core_op, BlockDiagonal)
+            )
+        )
+        for node in fg.toposort()
+    )
+
+    rewritten = rewrite_graph(bd, include=("canonicalize",))
+    utt.assert_equal_computations([rewritten], [A])
+
+
 def test_slogdet_blockdiag_rewrite():
     n_matrices = 10
     matrix_size = (5, 5)