jax.lax.broadcast — JAX documentation (original) (raw)

jax.lax.broadcast#

jax.lax.broadcast(operand, sizes, *, out_sharding=None)[source]#

Broadcasts an array, adding new leading dimensions

Parameters:

• operand (ArrayLike) – an array
• sizes (Sequence _[_int]) – a sequence of integers, giving the sizes of new leading dimensions to add to the front of the array.

Returns:

An array containing the result.

Return type:

Array

See also

jax.lax.broadcast_in_dim : add new dimensions at any location in the array shape.