jax-ml/jax · error · NotImplementedError

Unsupported device id type: {device_id_type}

Error message

Unsupported device id type: {device_id_type}

What it means

device_id_to_logical only supports DeviceIdType.MESH and DeviceIdType.LOGICAL; any other value falls through to NotImplementedError.

Source

Thrown at jax/_src/pallas/primitives.py:1293

    device_ids = tree_util.tree_leaves(device_id)
    mesh_strides: tuple[int, ...]
    if mesh_context is None:
      mesh_strides = ()
    else:
      mesh_strides = mesh_context.mesh_strides
    if len(device_ids) != len(mesh_strides):
      raise ValueError(
          "Number of device ids must match the number of mesh axes, but got"
          f" {len(device_ids)} ids for a {len(mesh_strides)}D mesh."
      )

    if not device_ids:
      # If there are no device ids, then it is purely local communication.
      return None, non_mesh_axes
    return sum(a * b for a, b in zip(device_ids, mesh_strides)), non_mesh_axes
  elif device_id_type is DeviceIdType.LOGICAL:
    return device_id, non_mesh_axes
  raise NotImplementedError(f"Unsupported device id type: {device_id_type}")


delay_p = jax_core.Primitive("delay")
delay_p.multiple_results = True


class DelayEffect(effects.Effect):
  pass
delay_effect = DelayEffect()
effects.control_flow_allowed_effects.add_type(DelayEffect)
pallas_core.kernel_local_effects.add_type(DelayEffect)


@delay_p.def_effectful_abstract_eval
def _delay_abstract_eval(nanos):
  del nanos
  return [], {delay_effect}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use DeviceIdType.MESH or DeviceIdType.LOGICAL explicitly
2. Update code that relies on removed enum values after a JAX upgrade

Example fix

// before
f(device_id_type=0)
// after
from jax._src.pallas.primitives import DeviceIdType
f(device_id_type=DeviceIdType.LOGICAL)

assert device_id_type in (DeviceIdType.MESH, DeviceIdType.LOGICAL), f"unsupported device_id_type: {device_id_type}"

def is_supported_device_id_type(t) -> bool:
    return t in (DeviceIdType.MESH, DeviceIdType.LOGICAL)

When it happens

Trigger: Passing an unknown/None/typo'd device_id_type value to APIs like semaphore_signal with remote devices.

Common situations: Custom or older code using a device_id_type enum value removed/renamed in this JAX version; passing device_id_type=None by accident.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• {axis} mixes JAX mesh and Pallas mesh grid axes
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• get not supported yet
• get not supported yet for block shape {b}

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