Repository: google/flax

Issue:
Imcompatibility with Flax Official ImageNet example with jax version >= 0.4.7
Hi, 

I was testing the [official flax example](https://github.com/google/flax/tree/main/examples/imagenet/) on Colab with jax and jaxlib version >= 0.4.7 on the colab pro+ environment with V100. After installing the requirements with `pip install -r requirements.txt` and with the following command `python main.py --workdir=./imagenet --config=configs/v100_x8.py`, the error is 

```
File "/content/FlaxImageNet/main.py", line 29, in <module>
import train
File "/content/FlaxImageNet/train.py", line 30, in <module>
from flax.training import checkpoints
File "/usr/local/lib/python3.10/dist-packages/flax/training/checkpoints.py", line 34, 
in <module>
from jax.experimental.global_device_array import GlobalDeviceArray
ModuleNotFoundError: No module named 'jax.experimental.global_device_array'
```

According to [this StackOverflow answer](https://stackoverflow.com/questions/76191911/no-module-named-jax-experimental-global-device-array-when-running-the-official/76192120#76192120), it seems that 'jax.experimental.global_device_array' is removed. 

Therefore, it would be great if one can fix the official example so that it works on newer version of jax. 
Unavailable to import checkpoints
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: All to its latest, also orbitax

Name: flax
Version: 0.6.9
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page: 
Author: 
Author-email: Flax team <flax-dev@google.com>
License: 
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: jax, msgpack, numpy, optax, orbax-checkpoint, PyYAML, rich, tensorstore, typing-extensions
Required-by: 
---
Name: jax
Version: 0.4.8
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: ml-dtypes, numpy, opt-einsum, scipy
Required-by: chex, diffrax, equinox, flax, optax, orbax, orbax-checkpoint, richmol
---
Name: jaxlib
Version: 0.4.7
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: ml-dtypes, numpy, scipy
Required-by: chex, optax, orbax, orbax-checkpoint
---
Name: orbax
Version: 0.1.7
Summary: Orbax
Home-page: 
Author: 
Author-email: Orbax Authors <orbax-dev@google.com>
License: 
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: absl-py, cached_property, etils, importlib_resources, jax, jaxlib, msgpack, nest_asyncio, numpy, pyyaml, tensorstore, typing_extensions

- Python version: 3.8


### Problem you have encountered:
When importing checkpoints, get the following error:
 """ 
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-1-0eac7b685376> in <module>
     11 config.update("jax_enable_x64", True)
     12 from flax import serialization
---> 13 from flax.training import checkpoints
     14 from jax import numpy as jnp
     15 import jax

/gpfs/cfel/group/cmi/common/psi4/psi4conda/lib//python3.8/site-packages/flax/training/checkpoints.py in <module>
     37 from jax import process_index
     38 from jax import sharding
---> 39 from jax.experimental.global_device_array import GlobalDeviceArray
     40 from jax.experimental.multihost_utils import sync_global_devices
     41 import orbax.checkpoint as orbax

ModuleNotFoundError: No module named 'jax.experimental.global_device_array'

"""

I guess it is a compatibility problem between jax and flax.

### What you expected to happen:

Usual importing




Discussion:

Can you try installing Flax from `main` and see if the issue is fixed?

```bash
pip install git+https://github.com/google/flax.git
```

We might need to create a new release.
@cgarciae Unluckily it does not make any difference.
Maybe try completely uninstalling `flax` first. The line on the error you are getting no longer exists on `main`:

https://github.com/google/flax/blob/74579637b4ef29087826665b82a9d0640730510f/flax/training/checkpoints.py#L37-L41

Resolving patch:
diff --git a/flax/version.py b/flax/version.py
--- a/flax/version.py
+++ b/flax/version.py
@@ -13,5 +13,5 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = "0.6.9"
+__version__ = "0.6.10"
 


Test patch: