jax

ff402188 - Expose reshard, auto_axes and explicit_axes via the `jax.sharding` endpoint graduating them from `jax.experimental`.