jax

76dec382 - Under pjit the `with mesh:` context will use `use_mesh(mesh): jit` instead of tracking separately using `resource_env`.

Under pjit the `with mesh:` context will use `use_mesh(mesh): jit` instead of tracking separately using `resource_env`. This would also make it easier to deprecate the `with mesh: pjit` path in the future from user code since the new path would be completely tested. This will also allow us to remove `resource_env` from JAX and the internal API access of `resource_env.physical_mesh` spread throughout codebases internally and externally. PiperOrigin-RevId: 735602187