Add concat_with_broadcast helper function

Co-authored-by: Ricardo <ricardo.vieira1994@gmail.com>

Author: jessegrabowski

pytensor/tensor/extra_ops.py

@@ -27,7 +27,7 @@
 from pytensor.scalar import upcast
 from pytensor.tensor import TensorLike, as_tensor_variable
 from pytensor.tensor import basic as ptb
-from pytensor.tensor.basic import alloc, second
+from pytensor.tensor.basic import alloc, join, second
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import abs as pt_abs
 from pytensor.tensor.math import all as pt_all
@@ -2018,6 +2018,42 @@ def broadcast_with_others(a, others):
     return brodacasted_vars
 
 
+def concat_with_broadcast(tensor_list, dim=0):
+    """
+    Concatenate a list of tensors, broadcasting the non-concatenated dimensions to align.
+    """
+    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 and not all_bcast[i])
+                or i == dim
+                or non_concat_shape[i] is not None
+            ):
+                continue
+            non_concat_shape[i] = sh
+
+    assert non_concat_shape.count(None) == 1
+
+    bcast_tensor_inputs = []
+    for tensor_inp in tensor_list:
+        # We modify the concat_axis in place, as we don't need the list anywhere else
+        non_concat_shape[dim] = tensor_inp.shape[dim]
+        bcast_tensor_inputs.append(broadcast_to(tensor_inp, non_concat_shape))
+
+    return join(dim, *bcast_tensor_inputs)
+
+
 __all__ = [
     "searchsorted",
     "cumsum",
@@ -2035,6 +2071,7 @@ def broadcast_with_others(a, others):
     "ravel_multi_index",
     "broadcast_shape",
     "broadcast_to",
+    "concat_with_broadcast",
     "geomspace",
     "logspace",
     "linspace",

tests/tensor/test_extra_ops.py

@@ -1333,3 +1333,32 @@ def test_space_ops(op, dtype, start, stop, num_samples, endpoint, axis):
         atol=1e-6 if config.floatX.endswith("64") else 1e-4,
         rtol=1e-6 if config.floatX.endswith("64") else 1e-4,
     )
+
+
+def test_concat_with_broadcast():
+    rng = np.random.default_rng()
+    a = pt.tensor("a", shape=(1, 3, 5))
+    b = pt.tensor("b", shape=(5, 3, 10))
+
+    c = pt.concat_with_broadcast([a, b], dim=-1)
+    fn = function([a, b], c, mode="FAST_COMPILE")
+    assert c.type.shape == (5, 3, 15)
+
+    a_val = rng.normal(size=(1, 3, 5))
+    b_val = rng.normal(size=(5, 3, 10))
+    c_val = fn(a_val, b_val)
+
+    # The result should be a tile + concat
+    np.testing.assert_allclose(c_val[:, :, :5], np.tile(a_val, (5, 1, 1)))
+    np.testing.assert_allclose(c_val[:, :, 5:], b_val)
+
+    # 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))