{"record":{"id":"d77eb6448cecf5ba","repo":"jax-ml/jax","slug":"index-type-must-be-i32","errorCode":null,"errorMessage":"`index_type` must be i32.","messagePattern":"`index_type` must be i32\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2540,"sourceCode":"  return _reduction_lowering(body, ctx, a, axes=axes)[0]\n\n\ntriton_lowering_rules[lax.reduce_max_p] = functools.partial(\n    _reduce_lowering, jnp.maximum\n)\ntriton_lowering_rules[lax.reduce_min_p] = functools.partial(\n    _reduce_lowering, jnp.minimum\n)\ntriton_lowering_rules[lax.reduce_sum_p] = functools.partial(\n    _reduce_lowering, jnp.add\n)\n\n\ndef _argreduce_lowering(\n    body, ctx: LoweringRuleContext, a, *, axes, index_dtype\n):\n  if index_dtype != jnp.int32:\n    raise ValueError(\"`index_type` must be i32.\")\n  if len(axes) != 1:\n    raise ValueError(\"`pallas` reduce operations only support one reduce axis.\")\n  [axis] = axes\n  [a_aval] = ctx.avals_in\n  index = _make_range(0, a_aval.shape[axis])\n  if len(a_aval.shape) > 1:\n    # Broadcast index across the non-reduced axes\n    for i in range(len(a_aval.shape)):\n      if i != axis:\n        index = _expand_dims(index, i)\n    index = _bcast_to(index, a_aval.shape)\n  ctx = ctx.replace(avals_in=[a_aval, a_aval.update(dtype=jnp.dtype(jnp.int32))])\n  _, indices = _reduction_lowering(body, ctx, (a, index), axes=axes)\n  return indices\n\n\ndef _reduce_argmax_combine(left, right):\n  value1, index1 = left","sourceCodeStart":2522,"sourceCodeEnd":2558,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2522-L2558","documentation":"Pallas arg-reduction lowering (argmax/argmin style ops) only supports int32 index outputs; the Triton index computation is hard-wired to i32. Passing index_dtype=int64/int16/uint32 raises ValueError before lowering.","triggerScenarios":"Calling jnp.argmax(x, index_dtype=jnp.int64) or lax.argmax(..., index_dtype=...) inside a Mosaic triton kernel.","commonSituations":"Code written for XLA backends that accept int64 indices (default in some ops) then run under pallas; kernels indexing large arrays where int64 was chosen to avoid overflow assumptions.","solutions":["Omit index_dtype so the op uses int32, or set index_dtype=jnp.int32 explicitly","Cast any downstream index arithmetic back to the desired dtype after the arg-reduction"],"exampleFix":"# before\nidx = jnp.argmax(x, index_dtype=jnp.int64)\n\n# after\nidx = jnp.argmax(x, index_dtype=jnp.int32)","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp\nassert index_dtype in (None, jnp.int32)","typeGuard":"def valid_index_dtype(dt) -> bool:\n    import jax.numpy as jnp\n    return dt is None or dt == jnp.int32","tryCatchPattern":null,"preventionTips":["Never pass index_dtype in pallas kernels; cast results afterwards"],"tags":["jax","pallas","triton","argmax","index-dtype"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}