jax

b707f0bd - [sharding_in_types] Error out when using `auto_axes` or `explicit_axes` API when there is no context mesh.

[sharding_in_types] Error out when using `auto_axes` or `explicit_axes` API when there is no context mesh. Those APIs don't support that right now anyways and they raise an ugly KeyError. Instead we raise a better error here. I have added a TODO to get the mesh from args so that computation follows data works but we can decide to do that in the future if a lot of users request that and don't want to use `use_mesh`. PiperOrigin-RevId: 730687231