jax

1629c6c7 - Make `jax.jit` work with vmap(..., spmd_axis_name) when there is no mesh context manager.

Make `jax.jit` work with vmap(..., spmd_axis_name) when there is no mesh context manager. This will only work if the input Array's sharding is a NamedSharding Fixes https://github.com/google/jax/issues/15886 PiperOrigin-RevId: 529758233