jax

21e1be33 - Don't call `get_cur_mesh_sharding` if sharding-in-types mode is not enabled