jax

840cf3f7 - [sharding_in_types] Add `pad_p` support to sharding_in_types to handle transpose to slice correctly.