jax-ml/jax · error · ValueError

`device_id`s must be an int32 value, but got

Error message

`device_id`s must be an int32 value, but got {aval.dtype}

What it means

In the semaphore_signal abstract eval, when remote signaling is requested every leaf of the device_id pytree must be an int32 value, matching the device addressing the hardware accepts. A non-int32 device id fails this check.

Solutions

1. Cast device ids to int32 before passing: tuple(int(i) for i in ids) via jnp.int32
2. Validate each leaf dtype of the device_id tree is int32

Example fix

// before
semaphore_signal(sem, 1, device_id=(0, dev_id))
// after
semaphore_signal(sem, 1, device_id=(0, jnp.int32(dev_id)))

ids = jax.tree.map(lambda x: x.astype(jnp.int32), device_id)
assert all(l.dtype == jnp.int32 for l in jax.tree.leaves(device_id))

When it happens

Trigger: Calling semaphore_signal(..., device_id=ids) where ids (or a dict entry) is int64/float rather than int32.

Common situations: Passing device ids obtained from jax.devices() indices or mesh computations as int64; mixed device_id dicts (MESH type) with non-int32 entries.

Related errors

• Must semaphores of the following types: . Got .
• Must signal an int32 value, but got
• Must wait an int32 value.
• mixes JAX mesh and Pallas mesh grid axes
• Cannot on a non-semaphore Ref

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

Appendix: source

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

    args_tree,
    device_id_type: DeviceIdType,
):
  (
      sem_aval,
      sem_transforms_avals,
      value_aval,
      device_id_aval,
      core_index_aval,
  ) = tree_util.tree_unflatten(args_tree, avals)
  check_sem_avals(sem_aval, sem_transforms_avals, "signal")
  if value_aval.dtype != jnp.dtype("int32"):
    raise ValueError(f"Must signal an int32 value, but got {value_aval.dtype}")
  effs: set[effects.Effect] = {sem_effect}
  if device_id_aval is not None:
    device_id_flat_avals = tree_util.tree_leaves(device_id_aval)
    for aval in device_id_flat_avals:
      if aval.dtype != jnp.dtype("int32"):
        raise ValueError(
            f"`device_id`s must be an int32 value, but got {aval.dtype}"
        )
    if device_id_type is DeviceIdType.MESH and isinstance(device_id_aval, dict):
      for k in device_id_aval:
        if not isinstance(k, tuple):
          k = (k,)
        for k_ in k:
          effs.add(jax_core.NamedAxisEffect(k_))
    else:
      effs.add(pallas_core.comms_effect)
  return [], effs


def _pp_device_id(device_id, context):
  if device_id is None:
    return pp.text("None")
  elif isinstance(device_id, dict):
    items = []

View on GitHub (pinned to 1e1c6a8fc0)