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.

[mvazquezg]

I managed to get it working by commenting out line 23 of model.py

#import jax.experimental.maps

and modifying the function at line 71:

def with_sharding_constraint(x, constraint):
    # Get the current mesh (if it exists)
    try:
        mesh = jax.sharding.thread_resources.env.physical_mesh
    except AttributeError:
        # If no active mesh is found, return x unchanged
        return x

    # Check if the mesh is empty
    if len(mesh.devices) == 0:
        return x
    else:
        # Apply the sharding constraint
        return pjit_sharding_constraint(x, constraint)