Use concat_with_broadcast in join split dot results

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -29,7 +29,6 @@
     constant,
     expand_dims,
     get_underlying_scalar_constant_value,
-    join,
     moveaxis,
     ones_like,
     register_infer_shape,
@@ -41,7 +40,7 @@
 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
+from pytensor.tensor.extra_ops import broadcast_arrays, concat_with_broadcast
 from pytensor.tensor.math import (
     Dot,
     Prod,
@@ -151,6 +150,7 @@ 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"""
@@ -174,29 +174,31 @@ def local_block_diag_dot_to_dot_block_diag(fgraph, node):
         ):
             continue
 
-        op = client.op
+        [blockdiag_result] = node.outputs
+        blockdiag_inputs = node.inputs
 
-        client_idx = client.inputs.index(node.outputs[0])
+        dot_op = client.op
 
+        client_idx = client.inputs.index(blockdiag_result)
         other_input = client.inputs[1 - client_idx]
-        components = node.inputs
 
         split_axis = -2 if client_idx == 0 else -1
         shape_idx = -1 if client_idx == 0 else -2
 
         other_dot_input_split = split(
             other_input,
-            splits_size=[component.shape[shape_idx] for component in components],
-            n_splits=len(components),
+            splits_size=[component.shape[shape_idx] for component in blockdiag_inputs],
+            n_splits=len(blockdiag_inputs),
             axis=split_axis,
         )
-        new_components = [
-            op(component, other_split)
+
+        split_dot_results = [
+            dot_op(component, other_split)
             if client_idx == 0
-            else op(other_split, component)
-            for component, other_split in zip(components, other_dot_input_split)
+            else dot_op(other_split, component)
+            for component, other_split in zip(blockdiag_inputs, other_dot_input_split)
         ]
-        new_output = join(split_axis, *new_components)
+        new_output = concat_with_broadcast(split_dot_results, dim=split_axis)
 
         copy_stack_trace(node.outputs[0], new_output)
         return {client.outputs[0]: new_output}