jax.lax.broadcast_in_dim

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

Wraps XLA’s BroadcastInDim operator.

Parameters

• operand (Any) –

• shape (Sequence[Union[int, Any]]) –

• broadcast_dimensions (Sequence[int]) –

Return type

Any