What causes "Cannot do int indexing on TPU" in jax-ml/jax?

Replace int indices with slices: use ref[i:i+1] (a zero-stride Slice) and squeeze the result instead of ref[i]

jax-ml/jax · error · ValueError

Cannot do int indexing on TPU

Error message

Cannot do int indexing on TPU

What it means

Raised by JAX's TPU Pallas/Mosaic lowering when a load/swap on a Ref uses an integer index in a transformation. TPU Pallas block references must be indexed with Slices (NumPy-style slice objects), not scalar ints, because lowering needs statically-known block extents. Any non-Slice index with a non-empty shape triggers this ValueError.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2271

      transforms,
      transforms_avals,
):
  if not transforms:
    prev_transforms, idx = [], NDIndexer.make_trivial_indexer(ref_aval.shape)
  else:
    if not isinstance(transforms[-1], NDIndexer):
      new_ref_aval = state.transform_type(transforms, ref_aval)
      assert isinstance(new_ref_aval, state.AbstractRef)
      idx = NDIndexer.make_trivial_indexer(new_ref_aval.shape)
      prev_transforms = transforms
    else:
      (*prev_transforms, idx) = transforms
      (*_, idx_aval) = transforms_avals
      if any(
          (not isinstance(a, primitives.Slice) and a.shape)
          for a in idx_aval.indices
      ):
        raise ValueError("Cannot do int indexing on TPU")
  return prev_transforms, idx


@register_lowering_rule(primitives.load_p, ensure_mlir_values=False)
def _load_lowering_rule(ctx: LoweringRuleContext, *args_flat, args_tree, **_):
  ref, transforms, mask, _ = args_tree.unflatten(args_flat)
  ref_aval, transforms_avals, _, _ = args_tree.unflatten(ctx.avals_in)
  prev_transforms, idx = _canonicalize_transforms_to_indexer(
      ref_aval, transforms, transforms_avals
  )
  if mask is not None:
    raise NotImplementedError
  if isinstance(ref_aval.memory_space, tpu_core.AccMemorySpace):
    raise ValueError(
        "Loading from an accumulator is not supported. Use `matmul_pop` "
        "instead, which will additionally zero out the accumulator."
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Replace int indices with slices: use ref[i:i+1] (a zero-stride Slice) and squeeze the result instead of ref[i]
Use jnp.take or do the scalar extraction outside the kernel on the host side
Check for squeeze/expand_dims patterns: keep block shapes non-scalar and slice them

Example fix

// before
val = ref[i]
// after
val = ref[i:i+1].squeeze(0)  # slice + squeeze instead of int indexing

Defensive patterns

Strategy: validation

Validate before calling

def _check_slice_indexing(idx):
    import numpy as np
    if isinstance(idx, tuple):
        for i in idx:
            if isinstance(i, (int, np.integer)) and not isinstance(i, slice):
                raise ValueError(f'int index {i!r} not allowed on TPU Pallas refs; use a slice')

Prevention

Always index TPU Pallas refs with slices or None/ellipsis, never bare ints
Centralize ref-access helpers in your kernel so indexing policy is checkable in one place

When it happens

Trigger: Calling pallas load/swap (or ref[...]) inside a TPU Pallas kernel with a plain int, e.g. ref[i] or ref[i, j], where the index aval is not a primitives.Slice. Produced by _canonicalize_transforms_to_indexer during lowering of load_p, PRNG key loads, or masked swaps.

Common situations: Porting a GPU (Triton-style) Pallas kernel to TPU; writing ref[0] to grab a scalar row; using dynamic_block_spec or manual indexing that JAX traces into int transforms.

Related errors

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