jax

af667199 - [sharding_in_types] Rename `.at[...].get(out_spec)` to `.at[...].get(out_sharding)`.