jax-ml/jax · error · ValueError

Unsupported device ID type: {device_id_type}

Error message

Unsupported device ID type: {device_id_type}

What it means

_device_id_to_logical received a DeviceIdType that is neither MESH nor LOGICAL. The interpreter only knows how to translate these two device ID representations into a logical ID.

Source

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

  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:
    raise ValueError(f"Unsupported device ID type: {device_id_type}")


def is_int(dtype):
  return jnp.issubdtype(dtype, jnp.integer)


def is_float(dtype):
  return jnp.issubdtype(dtype, jnp.floating)


@dataclasses.dataclass(frozen=True)
class Placeholder:
  """Placeholder for use in `JaxprEnv` below instead of storing a concrete value."""

  shape: tuple[int, ...]
  dtype: jnp.dtype

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade/downgrade so jax and jaxlib/pallas versions match (the enum should come from one consistent primitives module)
2. If writing a custom harness, use only DeviceIdType.MESH or DeviceIdType.LOGICAL
3. Check for stale imports of primitives from an old path

from jax._src.pallas import primitives
assert device_id_type in (primitives.DeviceIdType.MESH, primitives.DeviceIdType.LOGICAL), f'unsupported device id type: {device_id_type}'

def is_supported_device_id_type(t):
    return t in (primitives.DeviceIdType.MESH, primitives.DeviceIdType.LOGICAL)

When it happens

Trigger: Interpret-mode lowering encountering an unknown primitives.DeviceIdType enum value, typically from a version mismatch between jax and a plugin (jaxlib/pallas) or custom device-id plumbing.

Common situations: Mixing jax versions with experimental DeviceIdType variants; custom interpret harnesses constructing their own device_id_type values.

Related errors

• {axis} mixes JAX mesh and Pallas mesh grid axes
• {non_mesh_axes}
• Out-of-bounds read of ({device_id} {local_core_id} {memory_s
• masked load_p
• run_scoped_p with collective axes is not supported

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