look for _matmul in local_block_diag_dot_to_dot_block_diag

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -38,6 +38,7 @@
     zeros,
     zeros_like,
 )
+from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.extra_ops import broadcast_arrays
@@ -169,7 +170,7 @@ def local_block_diag_dot_to_dot_block_diag(fgraph, node):
                 isinstance(client.op, Dot)
                 and all(input.ndim == 2 for input in client.inputs)
             )
-            or client.op == _matrix_matrix_matmul
+            or client.op == _matmul
         ):
             continue