Look deeper for Dots

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -150,7 +150,6 @@ def local_0_dot_x(fgraph, node):
 
 
 @register_stabilize
-@register_specialize
 @node_rewriter([Blockwise])
 def local_block_diag_dot_to_dot_block_diag(fgraph, node):
     r"""
@@ -164,7 +163,9 @@ def local_block_diag_dot_to_dot_block_diag(fgraph, node):
         return
 
     # Check that the BlockDiagonal is an input to a Dot node:
-    for client in get_clients_at_depth(fgraph, node, depth=1):
+    for client in itertools.chain.from_iterable(
+        get_clients_at_depth(fgraph, node, depth=i) for i in [1, 2]
+    ):
         if not (
             (
                 isinstance(client.op, Dot)
@@ -179,7 +180,17 @@ def local_block_diag_dot_to_dot_block_diag(fgraph, node):
 
         dot_op = client.op
 
-        client_idx = client.inputs.index(blockdiag_result)
+        try:
+            client_idx = client.inputs.index(blockdiag_result)
+        except ValueError:
+            # If the blockdiag result is not an input to the dot, there is at least one Op between them (usually a
+            # DimShuffle). In this case, we need to figure out which of the inputs of the dot eventually leads to the
+            # blockdiag result.
+            for ancestor in client.inputs:
+                if ancestor.owner and blockdiag_result in ancestor.owner.inputs:
+                    client_idx = client.inputs.index(ancestor)
+                    break
+
         other_input = client.inputs[1 - client_idx]
 
         split_axis = -2 if client_idx == 0 else -1

tests/tensor/rewriting/test_math.py

@@ -4794,14 +4794,26 @@ def has_blockdiag(graph):
     fn = pytensor.function([a, b, c, d], out, mode=rewrite_mode)
     assert not has_blockdiag(fn.maker.fgraph.outputs[0])
 
+    n_dots_rewrite = sum(
+        isinstance(node.op, Dot | Dot22)
+        or (isinstance(node.op, Blockwise) and isinstance(node.op.core_op, Dot | Dot22))
+        for node in fn.maker.fgraph.apply_nodes
+    )
+    assert n_dots_rewrite == 3
+
     fn_expected = pytensor.function(
         [a, b, c, d],
         out,
         mode=Mode(linker="py", optimizer=None),
     )
     assert has_blockdiag(fn_expected.maker.fgraph.outputs[0])
 
-    # TODO: Count Dots
+    n_dots_no_rewrite = sum(
+        isinstance(node.op, Dot | Dot22)
+        or (isinstance(node.op, Blockwise) and isinstance(node.op.core_op, Dot | Dot22))
+        for node in fn_expected.maker.fgraph.apply_nodes
+    )
+    assert n_dots_no_rewrite == 1
 
     rng = np.random.default_rng()
     a_val = rng.normal(size=a.type.shape).astype(a.type.dtype)