jax.lax.broadcast_in_dim

jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)

Wraps XLA’s BroadcastInDim operator.

Parameters

• operand (Union[Array, ndarray, bool_, number, bool, int, float, complex]) – an array

• shape (Sequence[Union[int, Any]]) – the shape of the target array

• broadcast_dimensions (Sequence[int]) – to which dimension in the target shape each dimension of the operand shape corresponds to

Return type

Array

Returns

An array containing the result.

See also

jax.lax.broadcast : simpler interface to add new leading dimensions.