Tutorials
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
jit
vmap
pmap
jax.numpy.
broadcast_arrays
Like Numpy’s broadcast_arrays but doesn’t return views.