Test case with fake broadcast dim

Author: jessegrabowski

pytensor/tensor/extra_ops.py

@@ -2024,11 +2024,22 @@ def concat_with_broadcast(tensor_list, dim=0):
     """
     dim = dim if dim > 0 else tensor_list[0].ndim + dim
     non_concat_shape = [None] * tensor_list[0].ndim
+
+    # If all inputs have the same broadcastable dim, it's not really a broadcast
+    all_bcast = [
+        all(t.type.broadcastable[i] for t in tensor_list)
+        for i in range(tensor_list[0].ndim)
+    ]
+
     for tensor_inp in tensor_list:
         for i, (bcast, sh) in enumerate(
             zip(tensor_inp.type.broadcastable, tensor_inp.shape)
         ):
-            if bcast or i == dim or non_concat_shape[i] is not None:
+            if (
+                (bcast and not all_bcast[i])
+                or i == dim
+                or non_concat_shape[i] is not None
+            ):
                 continue
             non_concat_shape[i] = sh
 

tests/tensor/test_extra_ops.py

@@ -1349,3 +1349,14 @@ def test_concat_with_broadcast():
 
     c_val = fn(a_val, b_val)
     np.testing.assert_allclose(c_val[:, :, :5], np.tile(a_val, (5, 1, 1)))
+
+    # If a and b already conform, the result should be the same as a concatenation
+    a = pt.tensor("a", shape=(1, 1, 3, 5, 10))
+    b = pt.tensor("b", shape=(1, 1, 3, 2, 10))
+    c = pt.concatenate([a, b], axis=-2)
+    assert c.type.shape == (1, 1, 3, 7, 10)
+
+    fn = function([a, b], c, mode="FAST_COMPILE")
+    a_val, b_val = rng.normal(size=(1, 1, 3, 5, 10)), rng.normal(size=(1, 1, 3, 2, 10))
+    c_val = fn(a_val, b_val)
+    np.testing.assert_allclose(c_val, np.concatenate([a_val, b_val], axis=-2))