jax-ml/jax · critical · RuntimeError

Revisited block {output_ranges[i]} of output {i} in iteratio

Error message

Revisited block {output_ranges[i]} of output {i} in iteration {loop_idx}. The block was previously visited in iterations {past_output_ranges[past_idxs[0]][0]} through {past_output_ranges[past_idxs[-1]][0]} .

What it means

The TPU interpreter tracks which output blocks each grid iteration writes; writing the same output block in two different iterations is a kernel correctness bug (race/undefined ordering), so the interpreter raises RuntimeError with both the current and previously-covered iteration ranges.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:293

  shared_memory = _get_shared_memory()
  past_output_ranges = shared_memory.output_ranges[(device_id, local_core_id)]
  if not past_output_ranges:
    past_output_ranges.append((loop_idx, output_ranges))
    return token

  for i in range(len(output_ranges)):
    if output_ranges[i] is None:
      continue
    if past_output_ranges[-1][1][i] == output_ranges[i]:
      continue
    # TODO(jburnim): Do something constant time instead of linear here.
    past_idxs = [
        j
        for j, ors in enumerate(past_output_ranges)
        if ors[1][i] == output_ranges[i]
    ]
    if past_idxs:
      raise RuntimeError(
          f'Revisited block {output_ranges[i]} of output {i} in iteration '
          f'{loop_idx}. The block was previously visited in iterations '
          f'{past_output_ranges[past_idxs[0]][0]} through '
          f'{past_output_ranges[past_idxs[-1]][0]} .'
      )

  past_output_ranges.append((loop_idx, output_ranges))
  return token


@fail_on_exception
def _validate(token, device_id):
  device_id = int(device_id)

  shared_memory = _get_shared_memory()
  semaphores = shared_memory.get_sempahores_with_nonzero_count(device_id)
  if semaphores:
    sem, global_core_id = semaphores[0]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix index_fn so each grid iteration maps to a distinct output block (check multipliers like i * num_blocks)
2. Ensure output shape is divisible by block shape, or use masking instead of overlapping writes
3. Print/trace output_ranges across a tiny grid to find the colliding iterations
4. If overlap is intentional (e.g., reductions), restructure to accumulate in VMEM/SMEM and write once

Example fix

// before
out_specs=BlockSpec((B,), index_fn=lambda i, j: (j,))  # collides across i
// after
out_specs=BlockSpec((B,), index_fn=lambda i, j: (i * num_j + j,))

def index_map_unique(grid, index_fn, num_outputs):
    seen = set()
    for idx in itertools.product(*map(range, grid)):
        blocks = index_fn(*idx)
        key = tuple(map(tuple, blocks))
        if key in seen:
            return False, idx, key
        seen.add(key)
    return True, None, None

When it happens

Trigger: A pallas_call grid where the index_fn maps two different grid iterations to the same output block — e.g., an index map that ignores some grid index, or a block size vs grid shape mismatch causing overlapping writes.

Common situations: Reusing an index dimension in the index map (forgetting to multiply by block count); output shape not evenly divisible by block size so the tail block overlaps; refactored grid shapes without updating index maps.

Related errors

• Advanced indexers are not supported on TPU
• Dimensions with parallel semantics must form a prefix of the
• Non-trivial windowing is not supported for grid-free pallas_
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

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