jax-ml/jax · error · TypeError

Unsupported index type: {type(idx)}

Error message

Unsupported index type: {type(idx)}

What it means

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)}'.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:840

              start is not None and isinstance(start, int) and start % tile
          ) or (size is not None and isinstance(size, int) and size % tile):
            raise ValueError(
                f"Expected slice start ({start}) and slice size ({size})"
                f" to be divisible by the tile size ({tile})"
            )
          def _maybe_cdiv_with_cast(x, y):
            if x is None:
              return None
            if isinstance(x, jax.Array):
              # If x is an int32, we need to make sure y is an int32 to avoid
              # a dtype mismatch.
              y = jnp.array(y, x.dtype)
            return pallas_utils.cdiv(x, y)
          new_start = _maybe_cdiv_with_cast(start, tile)
          new_size = _maybe_cdiv_with_cast(size, tile)
          idxs_after_tiling.append(indexing.Slice(new_start, new_size))
        case _:
          raise TypeError(f"Unsupported index type: {type(idx)}")
    assert all(a % b == 0 for a, b in zip(untiled_shape, self.tiling))
    tiled_shape = [
        *(a // b for a, b in zip(untiled_shape, self.tiling)),
        *self.tiling,
    ]
    new_indexer = indexing.NDIndexer.from_indices_shape(
        indices=(*untiled_idxs, *idxs_after_tiling),
        shape=(*leading_shape, *tiled_shape)
    )
    return new_indexer, self

  def commute_reshape(
      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform
  ) -> tuple[state_types.ReshapeTransform, UntilingTransform]:
    if not transform.shape:
      raise NotImplementedError(
          "Commuting a `UntilingTransform` with a `ReshapeTransform` is not "
          "supported when the target shape has 0 dimensions"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the index to int, slice (stride 1), or a supported indexing.Slice before use
2. For boolean masks, compute integer indices (jnp.nonzero / arange) first and index with those
3. Check `type(idx)` in a debug pass and normalize the indexer pytree before the tiled call

Example fix

# before
x = ref[ref.shape[0] and mask]  # boolean/custom indexer -> TypeError

# after
rows = jnp.nonzero(mask)[0]  # or precomputed int indices
x = ref[rows]

from jax._src.pallas.mosaic import indexing
supported = (int, slice, indexing.Slice, indexing.NDIndexer)

def is_supported_index(idx) -> bool:
    from jax._src.pallas.mosaic import indexing
    return isinstance(idx, (int, slice, indexing.Slice))

null

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Unsupported range type: {type(r)}.
• Cannot do int indexing on TPU
• Can only load a single key per load.
• Dimension must be 0 for 1D iota.
• Integer indexing of refs that follows a non-trivial slice is

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