jax-ml/jax · error · RuntimeError

Computation failed on {self._failed_thread} with exception:\

Error message

Computation failed on {self._failed_thread} with exception:\n\n{failure_str}

What it means

A worker thread in the interpret-mode parallel simulation raised an exception; the main thread re-raises it wrapped as 'Computation failed on <thread>' with the original traceback. This is a propagation wrapper, not the root cause — the underlying exception text follows.

Source

Thrown at jax/_src/pallas/mosaic/interpret/shared_memory.py:882

    # it will call `set_failed` with `top_level=True`.
    if top_level:
      self.clean_up_barrier.wait()

  def check_failed(self):
    with self.lock:
      if self._failure is not None:
        # __cause__ information is lost when an exception from a callback goes
        # through XLA and back to Python, so we stuff the information into the
        # exception message.
        #
        # TODO(jburnim): The top-level exception (that the user sees) will not
        # always contain the original exception info/message.  It currently
        # depends on which raise (the raise here or the raise in the thread
        # that calls `set_failed`) happens first. (And maybe some other details
        # in XLA.)  We should figure out how to reliably propagate the
        # original exception.
        failure_str = "".join(traceback.format_exception(self._failure))
        raise RuntimeError(
            f"Computation failed on {self._failed_thread} with exception:\n\n{failure_str}"
        ) from None

  def append_semaphore_task(
      self,
      semaphore_id: int,
      global_core_id: int,
      task: SemaphoreTask,
  ):
    """Appends a task to be executed if the semaphore with the given sempahore ID is waiting to be signalled on the core with the given global core ID."""
    with self.lock:
      sem = self._unsafe_get_semaphore(semaphore_id)
    sem.enqueue_task(task, global_core_id)

  def get_random_virtual_device_id(self) -> int:
    # Virtual device IDs are needed for DMAs. Conceptually, each DMA runs on its
    # own, independent device. Representing this precisely would require vector
    # clocks to have sizes linear in the number of DMAs.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Read the traceback embedded after 'with exception:' and fix that root-cause error
2. Run the same kernel single-threaded (num_cores_or_threads=1) for cleaner stack traces
3. Reproduce with smaller grid/sizes to isolate the failing block index

try:
    run_kernel_interpreted(...)
except RuntimeError as e:
    if 'Computation failed on' in str(e):
        # the real traceback is embedded after 'with exception:'
        log.error('root cause:\n%s', e)  # fix the inner error, not this wrapper
    raise

When it happens

Trigger: Any exception raised inside a kernel body running on one of the interpreter's simulated threads; detected by ComputationManager when wait()/wrapper checks failed state.

Common situations: Bugs in user kernel code (shape errors, OOB, dtype mismatch) surfacing through the multithreaded interpreter; the real traceback is embedded in the message after 'exception:'.

Related errors

• Vector clock size ({self.vector_clock_size}) must be greater
• Logical shape {self.logical_shape} cannot be bigger than con
• Range {rnge} is entirely out of bounds for shape {self.shape
• Attempting to deallocate allocation with key `{key}` that is
• Attempting to get contents of allocation with key `{key}` th

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