jax-ml/jax · error · ValueError

Unsupported index: {v} of type {type(v)}

Error message

Unsupported index: {v} of type {type(v)}

What it means

The Mosaic GPU lowering could not convert a value into an index-typed MLIR value because its Python/JAX type is unrecognized. _as_index only accepts Python ints, index-capable ir.Values, and 0-d integer TypedNdArray literals. Anything else (floats, arrays with ndim>0, triton-style objects) hits this ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4471

  return arith_dialect.constant(ir.IntegerType.get_signless(64), v)


def _as_index(v: object) -> ir.Value:
  match v:
    case int():
      return arith_dialect.constant(ir.IndexType.get(), v)
    case ir.Value() if isinstance(v.type, ir.IndexType):
      return v
    case ir.Value() if isinstance(v.type, ir.IntegerType):
      return arith_dialect.index_cast(ir.IndexType.get(), v)
    case mgpu.FragmentedArray(layout=mgpu.WGSplatFragLayout()):
      return _as_index(v.registers.item())
    case jax_literals.TypedNdArray() if (
        np.issubdtype(v.dtype, np.integer) and v.ndim == 0
    ):
      return arith_dialect.constant(ir.IndexType.get(), int(v))
    case _:
      raise ValueError(f"Unsupported index: {v} of type {type(v)}")


def merge_indexers(
    indexers: Sequence[indexing.NDIndexer]) -> indexing.NDIndexer:
  """Merges multiple indexers into a single indexer.

  This function computes a new indexer such that applying the
  new indexer produces the same result as applying the sequence
  of input indexers in order from first-to-last.
  """
  if len(indexers) == 0:
    raise ValueError("Cannot merge empty list of indexers")
  if len(indexers) == 1:
    return indexers[0]
  root_shape = indexers[0].shape
  current_indices = [indexing.Slice(0, size, 1) for size in root_shape]
  removed_dimensions = set()
  for indexer in indexers:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure index functions return Python/numpy integer scalars (use int() or // integer division)
2. Check for accidental float division (/) in index computation; use //
3. If passing a jnp scalar, convert with int() first

Example fix

# before
def idx(block_idx): return block_idx[0] / 4  # float!
# after
def idx(block_idx): return block_idx[0] // 4

assert isinstance(v, (int, np.integer)) and getattr(v, 'ndim', 0) == 0, f"bad index {v!r}"

def is_valid_index(v) -> bool:
    import numpy as np
    return isinstance(v, (int, np.integer)) or (hasattr(v, 'ndim') and v.ndim == 0 and np.issubdtype(v.dtype, np.integer))

try:
    _as_index(v)
except ValueError as e:
    raise TypeError(f"convert index first: {e}") from e

When it happens

Trigger: Using a non-integer or non-scalar value (float, 1-d array, string, custom object) as an index/start/size inside a Mosaic GPU kernel's BlockSpec or slicing logic.

Common situations: Computing BlockSpec index functions that accidentally return numpy float arrays (e.g. numpy division producing floats) instead of int scalars; passing jnp arrays instead of python ints.

Related errors

• Barrier can only be indexed with integers or slices, got {id
• ref must be a reference
• Unsupported type: {x}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Can only store to references (got {x_ref}).

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