Getting Started
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
jit
vmap
pmap
jax.lax.
broadcast_in_dim
Wraps XLA’s BroadcastInDim operator.
operand (Any) –
Any
shape (Sequence[int]) –
Sequence
int
broadcast_dimensions (Sequence[int]) –