Add concat_with_broadcast helper function

Use new helper in xt.concat

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, axis=0):
+    """
+    Concatenate a list of tensors, broadcasting the non-concatenated dimensions to align.
+    """
+    if not tensor_list:
+        raise ValueError("Cannot concatenate an empty list of tensors.")
+
+    ndim = tensor_list[0].ndim
+    if not all(t.ndim == ndim for t in tensor_list):
+        raise TypeError(
+            "Only tensors with the same number of dimensions can be concatenated. "
+            f"Input ndims were: {[x.ndim for x in tensor_list]}"
+        )
+
+    axis = normalize_axis_index(axis=axis, ndim=ndim)
+    non_concat_shape = [1 if i != axis else None for i in range(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 == axis:
+                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[axis] = tensor_inp.shape[axis]
+        bcast_tensor_inputs.append(broadcast_to(tensor_inp, non_concat_shape))
+
+    return join(axis, *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",

pytensor/xtensor/rewriting/shape.py

@@ -2,8 +2,8 @@
 from pytensor.graph import node_rewriter
 from pytensor.tensor import (
     broadcast_to,
+    concat_with_broadcast,
     expand_dims,
-    join,
     moveaxis,
     specify_shape,
     squeeze,
@@ -74,28 +74,7 @@ def lower_concat(fgraph, node):
 
     # Convert input XTensors to Tensors and align batch dimensions
     tensor_inputs = [lower_aligned(inp, out_dims) for inp in node.inputs]
-
-    # Broadcast non-concatenated dimensions of each input
-    non_concat_shape = [None] * len(out_dims)
-    for tensor_inp in tensor_inputs:
-        # TODO: This is assuming the graph is correct and every non-concat dimension matches in shape at runtime
-        # I'm running this as "shape_unsafe" to simplify the logic / returned graph
-        for i, (bcast, sh) in enumerate(
-            zip(tensor_inp.type.broadcastable, tensor_inp.shape)
-        ):
-            if bcast or i == concat_axis 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_inputs:
-        # We modify the concat_axis in place, as we don't need the list anywhere else
-        non_concat_shape[concat_axis] = tensor_inp.shape[concat_axis]
-        bcast_tensor_inputs.append(broadcast_to(tensor_inp, non_concat_shape))
-
-    joined_tensor = join(concat_axis, *bcast_tensor_inputs)
+    joined_tensor = concat_with_broadcast(tensor_inputs, axis=concat_axis)
     new_out = xtensor_from_tensor(joined_tensor, dims=out_dims)
     return [new_out]
 

tests/tensor/test_extra_ops.py

@@ -1333,3 +1333,48 @@ 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], axis=2)
+    fn = function([a, b], c, mode="FAST_COMPILE")
+    assert c.type.shape == (5, 3, 15)
+
+    a_val = rng.normal(size=(1, 3, 5)).astype(config.floatX)
+    b_val = rng.normal(size=(5, 3, 10)).astype(config.floatX)
+    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.concat_with_broadcast([a, b], axis=-2)
+    assert c.type.shape == (1, 1, 3, 7, 10)
+
+    fn = function([a, b], c, mode="FAST_COMPILE")
+    a_val = rng.normal(size=(1, 1, 3, 5, 10)).astype(config.floatX)
+    b_val = rng.normal(size=(1, 1, 3, 2, 10)).astype(config.floatX)
+    c_val = fn(a_val, b_val)
+    np.testing.assert_allclose(c_val, np.concatenate([a_val, b_val], axis=-2))
+
+    c = pt.concat_with_broadcast([a], axis=0)
+    fn = function([a], c, mode="FAST_COMPILE")
+    np.testing.assert_allclose(fn(a_val), a_val)
+
+    with pytest.raises(ValueError, match="Cannot concatenate an empty list of tensors"):
+        pt.concat_with_broadcast([], axis=0)
+
+    with pytest.raises(
+        TypeError,
+        match="Only tensors with the same number of dimensions can be concatenated.",
+    ):
+        a = pt.tensor("a", shape=(1, 3, 5))
+        b = pt.tensor("b", shape=(3, 5))
+        pt.concat_with_broadcast([a, b], axis=1)