{"record":{"id":"8020342a82024caa","repo":"jax-ml/jax","slug":"advanced-indexers-are-not-supported-on-tpu","errorCode":null,"errorMessage":"Advanced indexers are not supported on TPU","messagePattern":"Advanced indexers are not supported on TPU","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":269,"sourceCode":"  # NOTE: We rely on the fact that `clean_up_barrier.wait()` will not raise.\n  # Otherwise, we could end up waiting on the barrier once here, and then again\n  # in the fail_on_exception wrapper -- so the barrier could complete without\n  # all devices having reached it.\n  shared_memory.clean_up_barrier.wait()\n  return token\n\n\n@fail_on_exception\ndef _check_for_revisiting(\n    token, device_id, local_core_id, loop_idx, output_blocks\n):\n  device_id = int(device_id)\n  local_core_id = int(local_core_id)\n  loop_idx = tuple(int(x) for x in loop_idx)\n  try:\n    output_blocks = jax.tree.map(int, output_blocks)\n  except:\n    raise ValueError('Advanced indexers are not supported on TPU')\n  output_ranges = [\n      interpret_utils.to_range(b) if b is not None else None\n      for b in output_blocks\n  ]\n\n  shared_memory = _get_shared_memory()\n  past_output_ranges = shared_memory.output_ranges[(device_id, local_core_id)]\n  if not past_output_ranges:\n    past_output_ranges.append((loop_idx, output_ranges))\n    return token\n\n  for i in range(len(output_ranges)):\n    if output_ranges[i] is None:\n      continue\n    if past_output_ranges[-1][1][i] == output_ranges[i]:\n      continue\n    # TODO(jburnim): Do something constant time instead of linear here.\n    past_idxs = [","sourceCodeStart":251,"sourceCodeEnd":287,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L251-L287","documentation":"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.","triggerScenarios":"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).","commonSituations":"Porting GPU kernels that use dynamic gather-based block indexing; index_fn using jnp operations on tracers instead of Python ints in interpret mode.","solutions":["Make index computations pure Python arithmetic on loop indices (ints), not traced jnp arrays","Precompute block index tables as concrete Python lists outside the kernel","Test dynamic-index kernels with the real compiler, not the interpreter"],"exampleFix":"// before\nindex_fn=lambda i, j: (idx_array[i, j],)  # traced array index\n// after\nindex_fn=lambda i, j: (i * num_blocks + j,)  # python int arithmetic","handlingStrategy":"validation","validationCode":"# ensure index maps return python ints before interpreting\nprobe = index_fn(*(0,) * index_fn.__code__.co_argcount)\nassert all(isinstance(i, int) for i in probe), 'index_fn must return ints'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Write index_fn using pure Python arithmetic","Precompute index tables as Python lists for interpret mode"],"tags":["jax","pallas","tpu","interpreter","advanced-indexing"],"backgroundTag":"unsupported-feature-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}