jax

3e742ff6 - Minor changes to sharding-in-types tutorial.

Minor changes to sharding-in-types tutorial. 1. Move imports of auto_axes and explicit_axes to their first usage. 2. Remove member import of set_mesh, since it's only used once. PiperOrigin-RevId: 789463709