jax

3e634d95 - [sharding_in_types] Add lax.transpose sharding propagation rule