jax-ml/jax · error · NotImplementedError

{non_mesh_axes}

Error message

{non_mesh_axes}

What it means

After converting a device-coordinate dict to mesh coordinates, some axes remained that are not part of either mesh — they could not be mapped into the logical program ID, so device_coords_to_logical_id raises NotImplementedError listing those axes.

Source

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

  device_id = []
  for axis in axis_names:
    if axis in physical_axis_dict:
      device_id.append(physical_axis_dict[axis])
    else:
      device_id.append(axis_indices[axis])
  non_mesh_axes = {
      k: v for k, v in physical_axis_dict.items() if k not in axis_names
  }
  return tuple(device_id), non_mesh_axes


def device_coords_to_logical_id(device_coords, axis_sizes, axis_indices):
  if isinstance(device_coords, dict):
    device_coords, non_mesh_axes = _device_id_dict_to_mesh(
        device_coords, axis_sizes, axis_indices
    )
    if non_mesh_axes:
      raise NotImplementedError(non_mesh_axes)
  if not isinstance(device_coords, tuple):
    device_coords = (device_coords,)
  assert len(device_coords) == len(axis_sizes)
  sizes = list(axis_sizes.values())
  ret = 0
  for i in range(len(device_coords)):
    ret += device_coords[i] * math.prod(sizes[i + 1 :])
  return ret


def _device_id_to_logical(device_id, device_id_type, axis_sizes, axis_indices):
  if device_id is None:
    return None
  if device_id_type == primitives.DeviceIdType.MESH:
    return device_coords_to_logical_id(device_id, axis_sizes, axis_indices)
  elif device_id_type == primitives.DeviceIdType.LOGICAL:
    return device_id
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove or correct the extraneous axes named in the error so all dict axes belong to the mesh/grid
2. Verify axis_sizes passed to interpret params matches the actual Mesh/Grid axis names
3. Regenerate device coordinates from the current mesh instead of hardcoding

known = set(axis_sizes)
extra = [a for k in device_coords if isinstance(k, tuple) for a in k if a not in known] or [k for k in device_coords if not isinstance(k, tuple) and k not in known]
assert not extra, f'unknown axes: {extra}'

When it happens

Trigger: Passing device coordinates (dict form) containing axis names that appear in neither the JAX mesh nor the Pallas grid axis_sizes; e.g. stale or misspelled axis names.

Common situations: Renaming mesh axes in one place but not another; leftover axes from a previous sharding config; hand-built device coordinate dicts in tests.

Related errors

• {axis} mixes JAX mesh and Pallas mesh grid axes
• Unsupported device ID type: {device_id_type}
• Unknown cluster axis {axis_name}, available axes: {[*axis_na
• No axis names are available. Make sure you are using `pl.cor
• Axis {axis_name} does not refer to a GPU mesh axis (availabl

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