jax

61d2016d - Merge pull request #17889 from jakevdp:jnp-axis-types