jax-ml/jax · error · NotImplementedError

Multiple core support not implemented.

Error message

Multiple core support not implemented.

What it means

core_index (multi-core signaling, e.g., TPU v5p/Vegeta cores) is not supported in the discharge/interpret path of semaphore_signal.

Source

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

    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(
    sem_or_view, value: int | jax_typing.Array = 1, *, decrement: bool = True
):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit core_index in interpret-mode runs
2. Run multicore signaling tests only on real TPU hardware via pallas_call

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

try:
    semaphore_signal(sem, 1, core_index=ci)
except NotImplementedError as e:
    if "Multiple core" in str(e):
        semaphore_signal(sem, 1)
    raise

When it happens

Trigger: Passing core_index to semaphore_signal while executing under interpret mode or a discharge-requiring path.

Common situations: Testing multicore TPU kernel semaphore logic locally in interpret mode.

Related errors

• Non-decrementing wait is not supported.
• Remote signal 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/fc1d68d827b5a99d. Report an issue: GitHub.