jax.experimental.maps deprecated. Grok-1 cannot be run any more.

[eliotli]

Tried to run the model but encountered this issue.

File "/home/jenkins/workspace/jax/releases/jax-release_jaxlib-v0.4.26-rocm61-bf16/grok-1/run.py", line 17, in
from model import LanguageModelConfig, TransformerConfig, QuantizedWeight8bit as QW8Bit
File "/home/jenkins/workspace/jax/releases/jax-release_jaxlib-v0.4.26-rocm61-bf16/grok-1/model.py", line 23, in
import jax.experimental.maps
ModuleNotFoundError: No module named 'jax.experimental.maps'

jax.experimental.maps is used in model.py. However, this module is deprecated in jax 0.4.31, which is a requirement for dm_haiku==0.0.12 installation. dm_haiku installation automatically upgrade the JAX version to 0.4.31.