Handle case with multiple clients

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -191,53 +191,54 @@ def check_for_block_diag(x):
         )
 
     # Check that the BlockDiagonal is an input to a Dot node:
-    clients = list(get_clients_at_depth(fgraph, node, depth=1))
-    if not clients or len(clients) > 1 or not isinstance(clients[0].op, Dot):
-        return
+    for client in get_clients_at_depth(fgraph, node, depth=1):
+        if not isinstance(client.op, Dot):
+            return
 
-    [dot_node] = clients
-    op = dot_node.op
-    x, y = dot_node.inputs
+        op = client.op
+        x, y = client.inputs
 
-    if not (check_for_block_diag(x) or check_for_block_diag(y)):
-        return None
+        if not (check_for_block_diag(x) or check_for_block_diag(y)):
+            return None
 
-    # Case 1: Only one input is BlockDiagonal. In this case, multiply all components of the block-diagonal with the
-    # non-block diagonal, and return a new block diagonal
-    if check_for_block_diag(x) and not check_for_block_diag(y):
-        components = x.owner.inputs
-        y_splits = split(
-            y,
-            splits_size=[component.shape[-1] for component in components],
-            n_splits=len(components),
-        )
-        new_components = [
-            op(component, y_split) for component, y_split in zip(components, y_splits)
-        ]
-        new_output = join(0, *new_components)
-
-    elif not check_for_block_diag(x) and check_for_block_diag(y):
-        components = y.owner.inputs
-        x_splits = split(
-            x,
-            splits_size=[component.shape[0] for component in components],
-            n_splits=len(components),
-            axis=1,
-        )
+        # Case 1: Only one input is BlockDiagonal. In this case, multiply all components of the block-diagonal with the
+        # non-block diagonal, and return a new block diagonal
+        if check_for_block_diag(x) and not check_for_block_diag(y):
+            components = x.owner.inputs
+            y_splits = split(
+                y,
+                splits_size=[component.shape[-1] for component in components],
+                n_splits=len(components),
+            )
+            new_components = [
+                op(component, y_split)
+                for component, y_split in zip(components, y_splits)
+            ]
+            new_output = join(0, *new_components)
+
+        elif not check_for_block_diag(x) and check_for_block_diag(y):
+            components = y.owner.inputs
+            x_splits = split(
+                x,
+                splits_size=[component.shape[0] for component in components],
+                n_splits=len(components),
+                axis=1,
+            )
 
-        new_components = [
-            op(x_split, component) for component, x_split in zip(components, x_splits)
-        ]
-        new_output = join(1, *new_components)
+            new_components = [
+                op(x_split, component)
+                for component, x_split in zip(components, x_splits)
+            ]
+            new_output = join(1, *new_components)
 
-    # Case 2: Both inputs are BlockDiagonal. Do nothing
-    else:
-        # TODO: If shapes are statically known and all components have equal shapes, we could rewrite
-        #  this case to block_diag(*[dot(comp_1, comp_2) for comp_1, comp_2 in zip(x.owner.inputs, y.owner.inputs)])
-        return None
+        # Case 2: Both inputs are BlockDiagonal. Do nothing
+        else:
+            # TODO: If shapes are statically known and all components have equal shapes, we could rewrite
+            #  this case to block_diag(*[dot(comp_1, comp_2) for comp_1, comp_2 in zip(x.owner.inputs, y.owner.inputs)])
+            return None
 
-    copy_stack_trace(node.outputs[0], new_output)
-    return {dot_node.outputs[0]: new_output}
+        copy_stack_trace(node.outputs[0], new_output)
+        return {client.outputs[0]: new_output}
 
 
 @register_canonicalize

tests/tensor/rewriting/test_math.py

@@ -4666,21 +4666,23 @@ def test_local_block_diag_dot_to_dot_block_diag(left_multiply):
     b = tensor("b", shape=(2, 4))
     c = tensor("c", shape=(4, 4))
     d = tensor("d", shape=(10, 10))
+    e = tensor("e", shape=(10, 10))
 
     x = pt.linalg.block_diag(a, b, c)
 
+    # Test multiple clients are all rewritten
     if left_multiply:
-        out = x @ d
+        out = [x @ d, x @ e]
     else:
-        out = d @ x
+        out = [d @ x, e @ x]
 
-    fn = pytensor.function([a, b, c, d], out, mode=rewrite_mode)
+    fn = pytensor.function([a, b, c, d, e], out, mode=rewrite_mode)
     assert not any(
         isinstance(node.op, BlockDiagonal) for node in fn.maker.fgraph.toposort()
     )
 
     fn_expected = pytensor.function(
-        [a, b, c, d],
+        [a, b, c, d, e],
         out,
         mode=rewrite_mode.excluding("local_block_diag_dot_to_dot_block_diag"),
     )
@@ -4690,10 +4692,11 @@ def test_local_block_diag_dot_to_dot_block_diag(left_multiply):
     b_val = rng.normal(size=b.type.shape).astype(b.type.dtype)
     c_val = rng.normal(size=c.type.shape).astype(c.type.dtype)
     d_val = rng.normal(size=d.type.shape).astype(d.type.dtype)
+    e_val = rng.normal(size=e.type.shape).astype(e.type.dtype)
 
     np.testing.assert_allclose(
-        fn(a_val, b_val, c_val, d_val),
-        fn_expected(a_val, b_val, c_val, d_val),
+        fn(a_val, b_val, c_val, d_val, e_val),
+        fn_expected(a_val, b_val, c_val, d_val, e_val),
         atol=1e-6 if config.floatX == "float32" else 1e-12,
         rtol=1e-6 if config.floatX == "float32" else 1e-12,
     )