jax

5634ca12 - Rename `out_shardings -> out_sharding` in the `auto_axes` API and `in_shardings -> in_sharding` in `explicit_axes` API.