jax.lax.broadcast_in_dim#
- jax.lax.broadcast_in_dim(operand, shape, broadcast_dimensions)[source]#
Wraps XLA’s BroadcastInDim operator.
- Parameters:
operand (
Union
[Array
,ndarray
,bool_
,number
,bool
,int
,float
,complex
]) – an arrayshape (
Sequence
[Union
[int
,Any
]]) – the shape of the target arraybroadcast_dimensions (
Sequence
[int
]) – to which dimension in the target shape each dimension of the operand shape corresponds to. That is, dimension i of the operand becomes dimension broadcast_dimensions[i] of the result.
- Return type:
- Returns:
An array containing the result.
See also
jax.lax.broadcast : simpler interface to add new leading dimensions.