jax-ml/jax · error · ValueError

Advanced indexers are not supported on TPU

Error message

Advanced indexers are not supported on TPU

What it means

While interpreting a Mosaic TPU kernel, output block indices are converted to ints. If any index is a tracer/array (i.e., computed with advanced/dynamic indexing), int() conversion fails and the interpreter reports that advanced indexers are unsupported on TPU.

Source

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

  # NOTE: We rely on the fact that `clean_up_barrier.wait()` will not raise.
  # Otherwise, we could end up waiting on the barrier once here, and then again
  # in the fail_on_exception wrapper -- so the barrier could complete without
  # all devices having reached it.
  shared_memory.clean_up_barrier.wait()
  return token


@fail_on_exception
def _check_for_revisiting(
    token, device_id, local_core_id, loop_idx, output_blocks
):
  device_id = int(device_id)
  local_core_id = int(local_core_id)
  loop_idx = tuple(int(x) for x in loop_idx)
  try:
    output_blocks = jax.tree.map(int, output_blocks)
  except:
    raise ValueError('Advanced indexers are not supported on TPU')
  output_ranges = [
      interpret_utils.to_range(b) if b is not None else None
      for b in output_blocks
  ]

  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 = [

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make index computations pure Python arithmetic on loop indices (ints), not traced jnp arrays
2. Precompute block index tables as concrete Python lists outside the kernel
3. Test dynamic-index kernels with the real compiler, not the interpreter

Example fix

// before
index_fn=lambda i, j: (idx_array[i, j],)  # traced array index
// after
index_fn=lambda i, j: (i * num_blocks + j,)  # python int arithmetic

# ensure index maps return python ints before interpreting
probe = index_fn(*(0,) * index_fn.__code__.co_argcount)
assert all(isinstance(i, int) for i in probe), 'index_fn must return ints'

When it happens

Trigger: A kernel whose out_specs/index computation uses advanced (array-based) indexing — e.g., index maps returning traced arrays from gather operations — run through the TPU Pallas interpreter, where _check_for_revisiting tries jax.tree.map(int, output_blocks).

Common situations: Porting GPU kernels that use dynamic gather-based block indexing; index_fn using jnp operations on tracers instead of Python ints in interpret mode.

Related errors

• Revisited block {output_ranges[i]} of output {i} in iteratio
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value
• run_scoped interpret rule does not support collective axes

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