ENH: Add a "broadcast" option to numpy.concatenate

Proposed new feature or change:

Mailing list post

Often, I've wanted to concatenate arrays with different ndims along a particular axis, broadcasting the other axes as needed. Others have sought this functionality as well:

• https://github.com/numpy/numpy/issues/2115
• https://stackoverflow.com/questions/56357047/concatenate-with-broadcast
• https://stackoverflow.com/questions/52733240/concatenate-1d-array-to-a-3d-array
• https://stackoverflow.com/questions/46700081/concat-two-arrays-of-different-dimensions-numpy
• https://stackoverflow.com/questions/55879664/how-to-concatenate-a-2d-array-into-every-3d-array
• https://stackoverflow.com/questions/63453495/combining-two-numpy-arrays

However, numpy.concatenate raises an error if the ndims don't match. For example:

import numpy as np

a = np.full([2], 0.1)
b = np.full([3, 2], 0.2)
c = np.full([5, 3, 2], 0.3)

arrays = [a, b, c]
axis = -1

try:
    np.concatenate(arrays, axis)
except ValueError as e:
    print(repr(e))

ValueError('all the input arrays must have same number of dimensions, but the array at index 0 has 1 dimension(s) and the array at index 1 has 2 dimension(s)')

It would be convenient to add to numpy.concatenate an optional boolean argument called broadcast that broadcasts the input arrays along the axes that are not the concatenation axis, before concatenating them. Its default value can be False, which is the current behavior.

Below is an example implementation:

def tuple_replace(tupl, index, item):
    return tupl[:index] + (item,) + tupl[index:][1:]

def broadcast_concat(arrays, axis):
    shape = np.broadcast_shapes(*(tuple_replace(a.shape, axis, 0) for a in arrays))
    bcast_arrays = [
        np.broadcast_to(a, tuple_replace(shape, axis, a.shape[axis])) for a in arrays
    ]
    return np.concatenate(bcast_arrays, axis)

output = broadcast_concat(arrays, axis)
assert output.shape[axis] == sum(a.shape[axis] for a in arrays)

If desired, I can submit a PR for this.

Start with the numpy.concatenate API reference and the linked mailing-list discussion, then compare the requested behavior with the current default behavior. Define the broadcast option's axis and shape semantics, including error cases, before implementation. Done means the optional behavior is specified while the default remains unchanged, with corresponding tests and documentation.