jax-ml/jax · error · NotImplementedError

{axis} mixes JAX mesh and Pallas mesh grid axes

Error message

{axis} mixes JAX mesh and Pallas mesh grid axes

What it means

A device-coordinate dict contained a tuple axis mixing JAX mesh axis names with Pallas mesh (grid) axis names. Mixed groupings are not supported when mapping device coordinates to a logical program ID.

Source

Thrown at jax/_src/pallas/mosaic/interpret/utils.py:160

  def __init__(self, initial_value: int):
    self.value = initial_value
    self.lock = threading.Lock()

  def get_next(self):
    with self.lock:
      result = self.value
      self.value += 1
    return result


# TODO(sharadmv): De-dup this w/ the impl in primitives.py.
def _device_id_dict_to_mesh(device_id_dict, axis_sizes, axis_indices):
  physical_axis_dict = {}
  axis_names = axis_sizes.keys()
  for axis, idx in device_id_dict.items():
    if isinstance(axis, tuple) and any(a in axis_names for a in axis):
      if not all(a in axis_names for a in axis):
        raise NotImplementedError(
            f"{axis} mixes JAX mesh and Pallas mesh grid axes"
        )
      axes_dimensions = [axis_sizes[name] for name in axis]
      for axis_index, axis_name in enumerate(axis):
        axis_size = axis_sizes[axis_name]
        inner_mesh_size = math.prod(axes_dimensions[axis_index + 1 :])
        minor_divisor = inner_mesh_size

        # Fast path for power of 2s
        if inner_mesh_size & (inner_mesh_size - 1) == 0:
          shift_len = (inner_mesh_size & -inner_mesh_size).bit_length() - 1
          partial_device_idx = idx >> shift_len
        else:
          partial_device_idx = idx // minor_divisor

        if axis_size & (axis_size - 1) == 0:
          device_idx = partial_device_idx & (axis_size - 1)
        else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the tuple axis so each tuple contains only JAX mesh axes or only Pallas grid axes
2. Rename axes so the two mesh namespaces don't overlap in one tuple key
3. Pass device_id as a plain coordinate tuple instead of a dict when possible

mesh_axes = set(axis_sizes)
for axis in device_id_dict:
    if isinstance(axis, tuple):
        assert all(a in mesh_axes for a in axis) or not any(a in mesh_axes for a in axis), 'mixed mesh/grid axis tuple'

When it happens

Trigger: Passing device_id as a dict whose keys are tuples where some names come from the JAX mesh (named_sharding axes) and others from the Pallas grid/mesh, when DeviceIdType.MESH is used in interpret mode.

Common situations: Using Mesh/JIT sharding with names that partially overlap Pallas grid axis names; constructing manual device coordinate dicts for multi-mesh interpret runs.

Related errors

• {non_mesh_axes}
• masked load_p
• run_scoped_p with collective axes is not supported
• Non-decrementing wait is not supported.
• {uninitialized_memory} + {str(dtype)}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9008df1105d034fb. Report an issue: GitHub.