{"record":{"id":"57797d7203381741","repo":"jax-ml/jax","slug":"cannot-do-int-indexing-on-tpu","errorCode":null,"errorMessage":"Cannot do int indexing on TPU","messagePattern":"Cannot do int indexing on TPU","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2271,"sourceCode":"      transforms,\n      transforms_avals,\n):\n  if not transforms:\n    prev_transforms, idx = [], NDIndexer.make_trivial_indexer(ref_aval.shape)\n  else:\n    if not isinstance(transforms[-1], NDIndexer):\n      new_ref_aval = state.transform_type(transforms, ref_aval)\n      assert isinstance(new_ref_aval, state.AbstractRef)\n      idx = NDIndexer.make_trivial_indexer(new_ref_aval.shape)\n      prev_transforms = transforms\n    else:\n      (*prev_transforms, idx) = transforms\n      (*_, idx_aval) = transforms_avals\n      if any(\n          (not isinstance(a, primitives.Slice) and a.shape)\n          for a in idx_aval.indices\n      ):\n        raise ValueError(\"Cannot do int indexing on TPU\")\n  return prev_transforms, idx\n\n\n@register_lowering_rule(primitives.load_p, ensure_mlir_values=False)\ndef _load_lowering_rule(ctx: LoweringRuleContext, *args_flat, args_tree, **_):\n  ref, transforms, mask, _ = args_tree.unflatten(args_flat)\n  ref_aval, transforms_avals, _, _ = args_tree.unflatten(ctx.avals_in)\n  prev_transforms, idx = _canonicalize_transforms_to_indexer(\n      ref_aval, transforms, transforms_avals\n  )\n  if mask is not None:\n    raise NotImplementedError\n  if isinstance(ref_aval.memory_space, tpu_core.AccMemorySpace):\n    raise ValueError(\n        \"Loading from an accumulator is not supported. Use `matmul_pop` \"\n        \"instead, which will additionally zero out the accumulator.\"\n    )\n","sourceCodeStart":2253,"sourceCodeEnd":2289,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2253-L2289","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nval = ref[i]\n// after\nval = ref[i:i+1].squeeze(0)  # slice + squeeze instead of int indexing","handlingStrategy":"validation","validationCode":"def _check_slice_indexing(idx):\n    import numpy as np\n    if isinstance(idx, tuple):\n        for i in idx:\n            if isinstance(i, (int, np.integer)) and not isinstance(i, slice):\n                raise ValueError(f'int index {i!r} not allowed on TPU Pallas refs; use a slice')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","pallas","tpu","indexing"],"backgroundTag":"pallas-tpu-slice-indexing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}