jax-ml/jax · error · ValueError

Must wait an int32 value.

Error message

Must wait an int32 value.

What it means

semaphore_wait requires its wait value to be int32; any other dtype raises this ValueError from the abstract eval.

Source

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

    sem_or_view: A Ref (or view) representing a semaphore.
    value: The target value that the semaphore should reach before unblocking.
    decrement: Whether to decrement the value of the semaphore after
      a successful wait.
  """
  ref, transforms = _get_ref_and_transforms(sem_or_view)
  value = jnp.asarray(value, dtype=jnp.int32)
  args = [ref, transforms, value, decrement]
  flat_args, args_tree = tree_util.tree_flatten(args)
  semaphore_wait_p.bind(*flat_args, args_tree=args_tree)

@semaphore_wait_p.def_effectful_abstract_eval
def _semaphore_wait_abstract_eval(*avals, args_tree):
  sem_aval, sem_transforms_avals, value_aval, _ = tree_util.tree_unflatten(
      args_tree, avals
  )
  check_sem_avals(sem_aval, sem_transforms_avals, "wait")
  if value_aval.dtype != jnp.dtype("int32"):
    raise ValueError("Must wait an int32 value.")
  return [], {sem_effect}

def _semaphore_wait_pp_eqn(eqn: jax_core.JaxprEqn,
                             context: jax_core.JaxprPpContext,
                             settings: jax_core.JaxprPpSettings):
  del settings
  invars = eqn.invars
  tree = eqn.params["args_tree"]
  (
      sem,
      sem_transforms,
      value,
      decrement,
  ) = tree_util.tree_unflatten(tree, invars)
  parts = [
      pp.text("semaphore_wait"),
  ]
  if decrement:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the value to int32: semaphore_wait(sem, value.astype(jnp.int32))
2. Construct the wait value as jnp.int32 from the beginning

Example fix

// before
semaphore_wait(sem, threshold)  # float32
// after
semaphore_wait(sem, jnp.int32(threshold))

assert value_aval.dtype == jnp.dtype('int32'), "semaphore_wait value must be int32"

When it happens

Trigger: Calling semaphore_wait(sem, value) where value is not int32 (e.g., float32 or int64).

Common situations: Passing computed thresholds in float; using Python ints promoted to int64 constants.

Related errors

• Must {name} semaphores of the following types: {allowed_sema
• Must signal an int32 value, but got {value_aval.dtype}
• `device_id`s must be an int32 value, but got {aval.dtype}
• Non-decrementing wait is not supported.
• Semaphore {sem_id} occurs as both fixed-id and internal.

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