Handle right-multiplication case

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -155,7 +155,7 @@ def local_0_dot_x(fgraph, node):
 @node_rewriter([Dot])
 def local_block_diag_dot_to_dot_block_diag(fgraph, node):
     r"""
-    Perform the rewrite ``dot(block_diag(A, B), C) -> block_diag(dot(A, C), dot(B, C))``
+    Perform the rewrite ``dot(block_diag(A, B), C) -> concat(dot(A, C), dot(B, C))``
 
     BlockDiag results in the creation of a matrix of shape ``(n1 * n2, m1 * m2)``. Because dot has complexity
     of approximately O(n^3), it's always better to perform two dot products on the smaller matrices, rather than
@@ -189,25 +189,18 @@ def check_for_block_diag(x):
         new_output = join(0, *new_components)
     elif not check_for_block_diag(x) and check_for_block_diag(y):
         components = y.owner.inputs
-        new_components = [op(x, component) for component in components]
-        new_output = join(0, *new_components)
-
-    # Case 2: Both inputs are BlockDiagonal. Here we can proceed only if the static shapes are known and identical. In
-    # that case, blockdiag(a,b) @ blockdiag(c, d) = blockdiag(a @ c, b @ d), but this is not true in the general case
-    elif any(shape is None for shape in (*x.type.shape, *y.type.shape)):
-        return None
-    elif x.ndim == y.ndim and all(
-        x_shape == y_shape for x_shape, y_shape in zip(x.type.shape, y.type.shape)
-    ):
-        x_components = x.owner.inputs
-        y_components = y.owner.inputs
+        x_splits = split(
+            x,
+            splits_size=[component.shape[0] for component in components],
+            n_splits=len(components),
+            axis=1,
+        )
 
-        if len(x_components) != len(y_components):
-            return None
+        new_components = [
+            op(x_split, component) for component, x_split in zip(components, x_splits)
+        ]
+        new_output = join(1, *new_components)
 
-        new_output = BlockDiagonal(len(x_components))(
-            *[op(x_comp, y_comp) for x_comp, y_comp in zip(x_components, y_components)]
-        )
     else:
         return None
 

tests/tensor/rewriting/test_math.py

@@ -4748,17 +4748,22 @@ def test_local_dot_to_mul(batched, a_shape, b_shape):
     )
 
 
-def test_local_block_diag_dot_to_dot_block_diag():
+@pytest.mark.parametrize("left_multiply", [True, False], ids=["left", "right"])
+def test_local_block_diag_dot_to_dot_block_diag(left_multiply):
     """
     Test that dot(block_diag(x, y,), z) is rewritten to concat(dot(x, z[:n]), dot(y, z[n:]))
     """
     a = tensor("a", shape=(4, 2))
     b = tensor("b", shape=(2, 4))
     c = tensor("c", shape=(4, 4))
-    d = tensor("d", shape=(10,))
+    d = tensor("d", shape=(10, 10))
 
     x = pt.linalg.block_diag(a, b, c)
-    out = x @ d
+
+    if left_multiply:
+        out = x @ d
+    else:
+        out = d @ x
 
     fn = pytensor.function([a, b, c, d], out)
     assert not any(