jax-ml/jax · error · ValueError

Dynamic indexing not supported in GPU interpret mode

Error message

Dynamic indexing not supported in GPU interpret mode

What it means

The GPU interpret mode only supports static (compile-time known) index expressions. If an NDIndexer transform contains dynamic (traced) values the interpreter cannot compute the read range on host and raises this ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:513

  else:
    uninit_array[tuple(slice(s) for s in ret.shape)] = ret
    return uninit_array


def _is_dynamic(indexer: indexing.NDIndexer) -> bool:
  return any(
      isinstance(idx, indexing.Slice)
      and (idx.is_dynamic_start or idx.is_dynamic_size)
      for idx in indexer.indices
  )


def _validate_transforms(transforms):
  for transform in transforms:
    match transform:
      case indexing.NDIndexer():
        if _is_dynamic(transform):
          raise ValueError(
              "Dynamic indexing not supported in GPU interpret mode"
          )
      case _:
        raise ValueError(f"Unsupported transform: {transform}")


def _get(
    token: jax.Array,
    mesh_location: memory.MeshLocation,
    thread: memory.Thread | None,
    allocation_key_as_array: jax.Array,
    transforms,
    block_indices=None,
    grid_loop_idx=None,
    clock=None,
    increment_clock: bool = True,
    source_info=None,
    input_name=None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all Block indices static: derive them from grid indices / python ints instead of traced values
2. Use static_unsafe or hoist the dynamic value to a Python scalar before constructing the Block
3. Test on device (non-interpret) if dynamic indexing is genuinely required, since interpret mode cannot emulate it
4. Refactor to express the varying index as a grid dimension instead

Example fix

# before
start = some_traced_scalar  # dynamic
blk = x[ds[start, bs], :]
# after
start = int(start_value)  # python int, static per interpretation step
blk = x[ds[start, bs], :]

from jax._src.pallas.mosaic_gpu.interpret.gpu_callbacks import _is_dynamic
assert not any(_is_dynamic(t) for t in blk.transforms), 'dynamic index'

def is_static_index(i) -> bool:
    return isinstance(i, (int,)) or not isinstance(i, jax.Array) or getattr(i, 'aval', None) is None

When it happens

Trigger: Constructing a Block whose dimension indices are functions of runtime values (e.g. loop-carried jnp scalars) and running the kernel in GPU interpret mode; any transform where _is_dynamic() detects non-static indices.

Common situations: Porting a TPU Pallas kernel that used dynamic indices to Mosaic GPU; indices derived from computed scalars rather than grid mappings; switching a kernel to interpret mode when indices depend on block arguments marked dynamic.

Related errors

• `thread_id` must be zero when allocating a buffer for all th
• `block_id` must be zero when allocating a buffer for all thr
• Out-of-bounds read of {allocation_key}: reading [{read_range
• Out-of-bounds block index {block_indices} for {allocation_ke
• Unsupported transform: {transform}

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