jax-ml/jax · error · ValueError

Expected an index-typed index

Error message

Expected an index-typed index

What it means

When indexing a memref with an ir.Value (a runtime/dynamic index), that value must have MLIR index type, not i32/i64. Mosaic's parse_indices raises ValueError because the index type determines how the memref.load/store indices are emitted.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:994

        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(start)
      slice_shape.append(stop - start)
      is_squeezed.append(False)
    elif isinstance(idx, DynamicSlice):
      if check_oob and (
          isinstance(idx.base, int) and idx.base + idx.length > bound
      ):
        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(idx.base)
      slice_shape.append(idx.length)
      is_squeezed.append(False)
    elif isinstance(idx, ir.Value):
      if not isinstance(idx.type, ir.IndexType):
        raise ValueError("Expected an index-typed index")
      base_indices.append(idx)
      slice_shape.append(1)
      is_squeezed.append(True)
    else:
      raise NotImplementedError(type(idx))
  assert len(base_indices) == len(slice_shape) == len(is_squeezed) == len(shape)
  return base_indices, slice_shape, is_squeezed


def commit_shared():
  nvvm.fence_proxy(
      nvvm.ProxyKind.async_shared, space=nvvm.SharedSpace.shared_cta
  )
  warpgroup_barrier()


def warpgroup_barrier_idx(sync: bool = True) -> ir.Value[ir.IntegerType]:
  # gpu.barrier() uses barrier number 0, and it would be unsafe to reuse it,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert first: idx = arith.index_cast(ir.IndexType.get(), i32_value)
2. Use utils.c(n, ir.IndexType.get()) for constants
3. Wrap dynamic values via arith.index_castui/index_cast depending on signedness before indexing

Example fix

# before
val = buf[i32_counter]
# after
idx = arith.index_cast(ir.IndexType.get(), i32_counter)
val = buf[idx]

if isinstance(idx, ir.Value) and not isinstance(idx.type, ir.IndexType):
    idx = arith.index_cast(ir.IndexType.get(), idx)

def is_index_typed(v) -> bool:
    return isinstance(v, ir.Value) and isinstance(v.type, ir.IndexType)

When it happens

Trigger: ref[arith.constant(5, i32)] or passing an i32 SSA value from a loop counter directly as an index, instead of converting with arith.index_cast to ir.IndexType.

Common situations: Building indices with integer arithmetic helpers that return i32/i64; mixing cuda/PTX-style i32 indices with MLIR memref indexing; copying test snippets that use ir.Value of the wrong width.

Related errors

• Only support memref.cast where the input and output types ar
• SubViewOp only supports a single tile transform.
• memref.cast transforms must have identical transforms for bo
• Unsupported memory space when lowering memref.cast: {memory_
• {op} has {name} of an unexpected type: {result}

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