jax-ml/jax · error · NotImplementedError

Remote signal not implemented.

Error message

Remote signal not implemented.

What it means

Signaling a semaphore on a remote device (device_id not None) is not implemented in the state-discharge (interpret-style) path of semaphore_signal.

Source

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

  ])
  if device_ids is not None:
    flat_device_ids = tree_util.tree_leaves(device_ids)
    if not flat_device_ids:
      return out
    out = pp.concat([out, pp.text(" "), _pp_device_id(device_ids, context)])
  return out

jax_core.pp_eqn_rules[semaphore_signal_p] = _semaphore_signal_pp_eqn


def _semaphore_signal_discharge_rule(ctx,
                                     *flat_args,
                                     args_tree,
                                     device_id_type):
  del device_id_type
  [ref, transforms, inc, device_id, core_index] = args_tree.unflatten(flat_args)
  if device_id is not None:
    raise NotImplementedError("Remote signal not implemented.")
  if core_index is not None:
    raise NotImplementedError("Multiple core support not implemented.")
  sem_value = _transform_semaphore(ref, transforms, ctx.in_avals[0])
  inc = inc.astype(pallas_core.SEMAPHORE_INTERPRET_DTYPE)
  _, new_sem_value = state_discharge.transform_swap_array(
      ref, transforms, sem_value + inc
  )
  return (new_sem_value,) + (None,) * (len(ctx.in_avals) - 1), ()
state_discharge.register_discharge_rule(semaphore_signal_p)(
    _semaphore_signal_discharge_rule
)


semaphore_wait_p = jax_core.Primitive('semaphore_wait')
semaphore_wait_p.multiple_results = True


def semaphore_wait(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop device_id when running under interpret mode / tests
2. Only use remote signaling in the real Pallas TPU execution path
3. Skip or mock remote-signal tests outside TPU

if INTERPRET_MODE and device_id is not None:
    raise unittest.SkipTest("remote signal unsupported in interpret mode")

try:
    semaphore_signal(sem, 1, device_id=did)
except NotImplementedError as e:
    if "Remote signal" in str(e):
        semaphore_signal(sem, 1)  # local fallback in tests
    raise

When it happens

Trigger: Using semaphore_signal with a device_id argument in a context that goes through discharge (interpret mode or non-Pallas lowering).

Common situations: Testing multicast/remote signaling code in interpret mode or on CPU where only local signaling is supported.

Related errors

• Non-decrementing wait is not supported.
• Multiple core support not implemented.
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• get not supported yet

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