{"record":{"id":"9f36f2901bf28550","repo":"jax-ml/jax","slug":"integer-indexing-of-refs-that-follows-a-non-trivia","errorCode":null,"errorMessage":"Integer indexing of refs that follows a non-trivial slice is not supported on SC","messagePattern":"Integer indexing of refs that follows a non-trivial slice is not supported on SC","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_lowering.py","lineNumber":115,"sourceCode":"  if not transforms or not isinstance(transforms[-1], indexing.NDIndexer):\n    tref_aval = state.transform_type(transforms, ref_aval)\n    assert isinstance(tref_aval, state.AbstractRef)\n    transforms.append(indexing.NDIndexer.make_trivial_indexer(tref_aval.shape))\n  *prev_transforms, indexer = transforms\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n  starts, sizes, strides, squeeze_dims, _ = tc_lowering._indexer_to_start_size_stride(\n      indexer, ref_block_shape, cast_to_index=True\n  )\n  for first_nontrivial_dim, s in enumerate(sizes):\n    if s != 1:\n      break\n  else:\n    first_nontrivial_dim = len(sizes)\n  if any(squeeze_dims[first_nontrivial_dim:]):\n    raise NotImplementedError(\n        \"Integer indexing of refs that follows a non-trivial slice is not\"\n        \" supported on SC\"\n    )\n  if not all(s == 1 for s in strides):\n    raise NotImplementedError(\n        \"Get only supports slices with stride 1, got {strides}\"\n    )\n\n  if (out_aval.ndim == 0) != (ref_memory_space is MemorySpace.SMEM):\n    message = \"Get only supports loading scalars from SMEM.\"\n    if ref_memory_space is MemorySpace.SMEM:\n      message += \" Trying to load an array of shape {out_aval.shape}.\"\n    elif ref_memory_space is MemorySpace.VMEM:\n      message += (\n          \" To load a scalar from VMEM, load an array first and then extract a\"\n          \" particular element, e.g. ``v = ref[pl.ds(idx, ...)]; v[0]``.\"\n      )\n    else:","sourceCodeStart":97,"sourceCodeEnd":133,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_lowering.py#L97-L133","documentation":"The SC Get lowering walks the ref index expression to find the first non-trivial (size != 1) dimension; any integer-index squeeze after that point cannot be lowered on SparseCore. Only leading trivial dims may be followed by scalar indexing.","triggerScenarios":"Indexing like ref[ds(...), i] — a real slice along one dimension followed by an integer index along a later dimension — in a SparseCore kernel read.","commonSituations":"Writing ref[block_slice, scalar_idx] (valid on other backends) when porting kernels to SC.","solutions":["Move integer indices before the non-trivial slice, or index the resulting array after the load: v = ref[slice]; v[i]","Reshape the ref so the integer-indexed dims come first"],"exampleFix":"# before\nv = ref[pl.ds(start, n), j]\n# after\nv = ref[pl.ds(start, n)]\njth = v[j]","handlingStrategy":"validation","validationCode":"# ensure no squeeze dims after first nontrivial slice dim\nassert not any(squeeze_dims[first_nontrivial_dim:])","typeGuard":"null","tryCatchPattern":"null","preventionTips":["Load slices first, then index the resulting JAX array"],"tags":["jax","pallas","tpu","sparsecore","indexing"],"backgroundTag":"unsupported-indexing-pattern","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}