{"record":{"id":"7aef54a4fc312a1b","repo":"jax-ml/jax","slug":"unsupported-index-type-type-idx","errorCode":null,"errorMessage":"Unsupported index type: {type(idx)}","messagePattern":"Unsupported index type: (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":840,"sourceCode":"              start is not None and isinstance(start, int) and start % tile\n          ) or (size is not None and isinstance(size, int) and size % tile):\n            raise ValueError(\n                f\"Expected slice start ({start}) and slice size ({size})\"\n                f\" to be divisible by the tile size ({tile})\"\n            )\n          def _maybe_cdiv_with_cast(x, y):\n            if x is None:\n              return None\n            if isinstance(x, jax.Array):\n              # If x is an int32, we need to make sure y is an int32 to avoid\n              # a dtype mismatch.\n              y = jnp.array(y, x.dtype)\n            return pallas_utils.cdiv(x, y)\n          new_start = _maybe_cdiv_with_cast(start, tile)\n          new_size = _maybe_cdiv_with_cast(size, tile)\n          idxs_after_tiling.append(indexing.Slice(new_start, new_size))\n        case _:\n          raise TypeError(f\"Unsupported index type: {type(idx)}\")\n    assert all(a % b == 0 for a, b in zip(untiled_shape, self.tiling))\n    tiled_shape = [\n        *(a // b for a, b in zip(untiled_shape, self.tiling)),\n        *self.tiling,\n    ]\n    new_indexer = indexing.NDIndexer.from_indices_shape(\n        indices=(*untiled_idxs, *idxs_after_tiling),\n        shape=(*leading_shape, *tiled_shape)\n    )\n    return new_indexer, self\n\n  def commute_reshape(\n      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform\n  ) -> tuple[state_types.ReshapeTransform, UntilingTransform]:\n    if not transform.shape:\n      raise NotImplementedError(\n          \"Commuting a `UntilingTransform` with a `ReshapeTransform` is not \"\n          \"supported when the target shape has 0 dimensions\"","sourceCodeStart":822,"sourceCodeEnd":858,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L822-L858","documentation":"TilingTransform.commute_ndindexer handles only ints, slice/Slice, and similar supported index types in its match statement. Any other indexer object (unknown Index type, fancy index, etc.) raises TypeError 'Unsupported index type: {type(idx)}'.","triggerScenarios":"Indexing a tiled ref with an unsupported indexer object, e.g. a boolean mask, a custom indexing class, or a non-standard Index replacement inside a Pallas kernel.","commonSituations":"Porting NumPy-style fancy/boolean indexing into Pallas kernels; custom Index subclasses from newer indexing APIs reaching Mosaic lowering; pytrees where a leaf index object isn't recognized.","solutions":["Convert the index to int, slice (stride 1), or a supported indexing.Slice before use","For boolean masks, compute integer indices (jnp.nonzero / arange) first and index with those","Check `type(idx)` in a debug pass and normalize the indexer pytree before the tiled call"],"exampleFix":"# before\nx = ref[ref.shape[0] and mask]  # boolean/custom indexer -> TypeError\n\n# after\nrows = jnp.nonzero(mask)[0]  # or precomputed int indices\nx = ref[rows]","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas.mosaic import indexing\nsupported = (int, slice, indexing.Slice, indexing.NDIndexer)","typeGuard":"def is_supported_index(idx) -> bool:\n    from jax._src.pallas.mosaic import indexing\n    return isinstance(idx, (int, slice, indexing.Slice))","tryCatchPattern":"null","preventionTips":["Normalize custom/boolean indexers to ints or slices before tiled kernels","Avoid fancy indexing inside Pallas kernels on tiled refs"],"tags":["jax","pallas","indexing","type-dispatch","tiling"],"backgroundTag":"unsupported-index-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}