Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
jit
vmap
pmap
jax.numpy.
msort
Return a copy of an array sorted along the first axis.
LAX-backend implementation of msort().
msort()
Original docstring below.
a (array_like) – Array to be sorted.
sorted_array – Array of the same type and shape as a.
ndarray