jax

353363ae - Move jax._src.shard_map to its own BUILD rule

Move jax._src.shard_map to its own BUILD rule Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, prevents use of internal APIs, and leads to improved build and iteration times. PiperOrigin-RevId: 774532216