{"record":{"id":"c5f97fbf39e75539","repo":"jax-ml/jax","slug":"unsupported-combination-of-input-dtype-x-aval-dt","errorCode":null,"errorMessage":"Unsupported combination of input dtype ({x_aval.dtype}) and index_dtype ({index_dtype}) for reduce_index","messagePattern":"Unsupported combination of input dtype \\((.+?)\\) and index_dtype \\((.+?)\\) for reduce_index","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3563,"sourceCode":"  (aval_out,) = ctx.avals_out\n  if jnp.issubdtype(aval_out.dtype, jnp.signedinteger):\n    return arith.minsi(x, y)\n  elif jnp.issubdtype(aval_out.dtype, jnp.unsignedinteger):\n    return arith.minui(x, y)\n  elif jnp.issubdtype(aval_out.dtype, jnp.floating):\n    return arith.minimumf(x, y)\n  raise NotImplementedError(aval_out.dtype)\n\ndef _reduce_index_helper(\n    ctx: LoweringRuleContext, x, axes, index_dtype, reduction_kind):\n  (x_aval,) = ctx.avals_in\n  (out_aval,) = ctx.avals_out\n  if (x_aval.dtype, index_dtype) not in (\n      (jnp.float32, jnp.int32),\n      (jnp.bfloat16, jnp.int16),\n      (jnp.bfloat16, jnp.int32),\n  ):\n    raise NotImplementedError(\n        f\"Unsupported combination of input dtype ({x_aval.dtype}) and\"\n        f\" index_dtype ({index_dtype}) for reduce_index\"\n    )\n  if len(axes) != 1:\n    raise NotImplementedError(\"Only single axis reduction supported\")\n\n  axis = axes[0]\n  # TODO(b/460843515): Support 1D inputs in Mosaic.\n  is_1d = len(x_aval.shape) == 1\n  if is_1d:\n    x = vector.shape_cast(\n        ctx.aval_to_ir_type(\n            jax_core.ShapedArray((1, *x_aval.shape), x_aval.dtype)\n        ),\n        x,\n    )\n    axis += 1\n    out_shape = (1, *out_aval.shape)","sourceCodeStart":3545,"sourceCodeEnd":3581,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3545-L3581","documentation":"The reduce_index lowering (argmin/argmax-style reductions that return indices) only accepts the (input dtype, index dtype) pairs (float32,int32), (bfloat16,int16), (bfloat16,int32). Any other combination raises NotImplementedError at lowering time.","triggerScenarios":"Calling jnp.argmin/jnp.argmax (or lax.argmin/argmax) inside a Pallas kernel where the operand is not float32/bfloat16 or the index_dtype parameter is not int16/int32 (e.g. default uint16/uint32 index dtype on some paths, or float16 input).","commonSituations":"Using argmax over float16 logits inside a kernel (common in inference kernels); lax.argmax with index_dtype=jnp.uint16; mixing bfloat16 with an int64 index request.","solutions":["Cast the operand to jnp.float32 (or bfloat16) before argmax/argmin","Pass index_dtype explicitly: lax.argmax(x, axis, index_dtype=jnp.int32)","Compute indices outside the kernel","Upgrade JAX — reduce_index dtype coverage has been expanding"],"exampleFix":"// before\ni = lax.argmax(x, axis=0, index_dtype=jnp.uint16)\n// after\ni = lax.argmax(x.astype(jnp.float32), axis=0, index_dtype=jnp.int32)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nSUPPORTED = {(jnp.float32, jnp.int32), (jnp.bfloat16, jnp.int16), (jnp.bfloat16, jnp.int32)}\ndef reduce_index_ok(x_dtype, idx_dtype):\n    return any(jnp.dtype(x_dtype) == a and jnp.dtype(idx_dtype) == b for a, b in SUPPORTED)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass index_dtype=jnp.int32 to argmax/argmin in kernels","Cast operands to float32 or bfloat16 before index reductions"],"tags":["jax","pallas","tpu","argmax","dtype"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}