jax.lax.broadcast_in_dim

Warning

This page was created from a pull request (#9655).

jax.lax.broadcast_in_dim¶

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

Wraps XLA’s BroadcastInDim operator.

Parameters

  • operand (Any) – 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

Any

Returns

An array containing the result.

See also

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

previous

jax.lax.broadcasted_iota

next

jax.lax.cbrt