array_api_extra.broadcast_shapes

array_api_extra.broadcast_shapes(*shapes)

Compute the shape of the broadcasted arrays.

Duplicates numpy.broadcast_shapes(), with additional support for None and NaN sizes.

This is equivalent to xp.broadcast_arrays(arr1, arr2, ...)[0].shape without needing to worry about the backend potentially deep copying the arrays.

Parameters:

*shapes (tuple[int | None, ...]) – Shapes of the arrays to broadcast.

Returns:

The shape of the broadcasted arrays.

Return type:

tuple[int | None, …]

See also

numpy.broadcast_shapes

Equivalent NumPy function.

array_api.broadcast_arrays

Function to broadcast actual arrays.

Notes

This function accepts the Array API’s None for unknown sizes, as well as Dask’s non-standard math.nan. Regardless of input, the output always contains None for unknown sizes.

Examples

>>> import array_api_extra as xpx
>>> xpx.broadcast_shapes((2, 3), (2, 1))
(2, 3)
>>> xpx.broadcast_shapes((4, 2, 3), (2, 1), (1, 3))
(4, 2, 3)