jax.lax.broadcast_in_dim

jax.lax.broadcast_in_dim#

jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)[source]#

Wraps XLA’s BroadcastInDim operator.

Parameters:
Return type:

Array

Returns:

An array containing the result.

See also

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