{"record":{"id":"b700feaf2c7593ce","repo":"jax-ml/jax","slug":"expected-an-index-typed-index","errorCode":null,"errorMessage":"Expected an index-typed index","messagePattern":"Expected an index-typed index","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":994,"sourceCode":"        raise IndexError(\n            f\"Slice {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(start)\n      slice_shape.append(stop - start)\n      is_squeezed.append(False)\n    elif isinstance(idx, DynamicSlice):\n      if check_oob and (\n          isinstance(idx.base, int) and idx.base + idx.length > bound\n      ):\n        raise IndexError(\n            f\"Slice {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(idx.base)\n      slice_shape.append(idx.length)\n      is_squeezed.append(False)\n    elif isinstance(idx, ir.Value):\n      if not isinstance(idx.type, ir.IndexType):\n        raise ValueError(\"Expected an index-typed index\")\n      base_indices.append(idx)\n      slice_shape.append(1)\n      is_squeezed.append(True)\n    else:\n      raise NotImplementedError(type(idx))\n  assert len(base_indices) == len(slice_shape) == len(is_squeezed) == len(shape)\n  return base_indices, slice_shape, is_squeezed\n\n\ndef commit_shared():\n  nvvm.fence_proxy(\n      nvvm.ProxyKind.async_shared, space=nvvm.SharedSpace.shared_cta\n  )\n  warpgroup_barrier()\n\n\ndef warpgroup_barrier_idx(sync: bool = True) -> ir.Value[ir.IntegerType]:\n  # gpu.barrier() uses barrier number 0, and it would be unsafe to reuse it,","sourceCodeStart":976,"sourceCodeEnd":1012,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L976-L1012","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Convert first: idx = arith.index_cast(ir.IndexType.get(), i32_value)","Use utils.c(n, ir.IndexType.get()) for constants","Wrap dynamic values via arith.index_castui/index_cast depending on signedness before indexing"],"exampleFix":"# before\nval = buf[i32_counter]\n# after\nidx = arith.index_cast(ir.IndexType.get(), i32_counter)\nval = buf[idx]","handlingStrategy":"type-guard","validationCode":"if isinstance(idx, ir.Value) and not isinstance(idx.type, ir.IndexType):\n    idx = arith.index_cast(ir.IndexType.get(), idx)","typeGuard":"def is_index_typed(v) -> bool:\n    return isinstance(v, ir.Value) and isinstance(v.type, ir.IndexType)","tryCatchPattern":null,"preventionTips":["Convert all dynamic counters to index type at the point of creation","Use utils.c(n, ir.IndexType.get()) for constant indices"],"tags":["jax","mosaic-gpu","memref","mlir","type-mismatch"],"backgroundTag":"mlir-index-type-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}