Search code examples
pythonnumpy

Correct way to find dimension after broadcasting in numpy

Suppose I am given a few numpy arrays, say a, b and c, which are assumed to be broadcastable. Is there a standard or othwerwise an elegant way to find the array shape after broadcasting? Of course, something like (a+b+c).shape would work, but is very inefficient if I am only interested in the shape of the result.

Solution

  • You can use broadcast_shapes. Like this

    result_shape = np.broadcast_shapes(a.shape, b.shape, c.shape)

    Here is the documentation page: https://numpy.org/doc/stable/reference/generated/numpy.broadcast_shapes.html

    Good luck!