imagenet example needs a refreshed requirements.txt

it seems like the flax examples could use a version bump?

### System information
- OS Platform and Distribution: Ubuntu 22.04.4 LTS
- Flax, jax, jaxlib versions:

```
pip show flax jax jaxlib
Name: flax
Version: 0.6.5
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page: https://github.com/google/flax
Author: Flax team
Author-email: flax-dev@google.com
License:
Location: ~/src/python/flax/examples/imagenet/.venv/lib/python3.10/site-packages
Requires: jax, matplotlib, msgpack, numpy, optax, orbax, PyYAML, rich, tensorstore, typing-extensions
Required-by: clu
---
Name: jax
Version: 0.4.28
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: ~/src/python/flax/examples/imagenet/.venv/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, opt-einsum, scipy
Required-by: chex, clu, flax, optax, orbax-checkpoint
---
Name: jaxlib
Version: 0.4.28+cuda12.cudnn89
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: ~/src/python/flax/examples/imagenet/.venv/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, scipy
Required-by: chex, clu, optax, orbax-checkpoint

```
- Python version:
```
python3 --version
Python 3.10.12
```
- GPU/TPU model and memory: RTX 4070 12GB
- CUDA version (if applicable):
```
nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Jun_13_19:16:58_PDT_2023
Cuda compilation tools, release 12.2, V12.2.91
Build cuda_12.2.r12.2/compiler.32965470_0
```

### Problem you have encountered:
imagenet with venv:
If I create a venv and install the requirements.txt, :
```
File "~/src/python/flax/examples/imagenet/.venv/lib/python3.10/site-packages/flax/configurations.py", line 74, in
flax_filter_frames = define_bool_state(
File "~/src/python/flax/examples/imagenet/.venv/lib/python3.10/site-packages/flax/configurations.py", line 40, in define_bool_state
return jax_config.define_bool_state('flax_' + name, default, help)
AttributeError: 'Config' object has no attribute 'define_bool_state'
```
My pip list flax version reports 0.6.5 and the requirements.txt specifies 0.6.5

### What you expected to happen:
imagenet trains

### Logs, error messages, etc:
see above

### Steps to reproduce:
create a fresh and up to date ubuntu install with an nvidia card, clone flax, and run the requirements.txt