{"record":{"id":"0cba2784b980c960","repo":"jax-ml/jax","slug":"index-dtype-must-be-an-integer-type-but-got","errorCode":null,"errorMessage":"index_dtype must be an integer type, but got {}","messagePattern":"index_dtype must be an integer type, but got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8672,"sourceCode":"batching.defreducer(reduce_min_p)\n\ndef _argminmax_shape_rule(operand, *, axes, index_dtype):\n  axis, = axes\n  if not (0 <= axis < len(operand.shape)):\n    raise ValueError(f\"Invalid axis {axis} for operand shape {operand.shape}\")\n  if operand.shape[axis] < 1:\n    raise ValueError(\"argmin and argmax require non-empty reduced dimension. \"\n                     f\"operand.shape={operand.shape} {axis=}\")\n  return util.tuple_delete(operand.shape, axis)\n\ndef _argminmax_sharding_rule(operand, *, axes, index_dtype):\n  axis, = axes\n  return operand.sharding.update(spec=\n      util.tuple_delete(operand.sharding.spec, axis))\n\ndef _argminmax_dtype_rule(operand, *, axes, index_dtype):\n  if not dtypes.issubdtype(index_dtype, np.integer):\n    raise TypeError(\"index_dtype must be an integer type, but got {}\"\n                    .format(dtype_to_string(index_dtype)))\n  return index_dtype\n\nclass _ArgMinMaxReducer:\n\n  def __init__(self, value_comparator: Callable[[Any, Any], Any]):\n    self._value_comparator = value_comparator\n\n  def __repr__(self):\n    # Override the repr so that the metadata attached to the lowered op does not\n    # contain unstable function ids. This plays more nicely with computation\n    # fingerprint calculation in the compilation cache.\n    return f'_ArgMinMaxReducer({self._value_comparator.__name__})'\n\n  def __call__(self, op_val_index, acc_val_index):\n    op_val, op_index = op_val_index\n    acc_val, acc_index = acc_val_index\n    # Pick op_val if Lt (for argmin) or if NaN","sourceCodeStart":8654,"sourceCodeEnd":8690,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8654-L8690","documentation":"argmin/argmax's index_dtype parameter must be a NumPy integer type (e.g. int32/int64), but something else (float, bool, or a non-dtype) was supplied. This is a TypeError raised during dtype-rule evaluation.","triggerScenarios":"lax.argmin(x, axis, index_dtype=jnp.float32), passing index_dtype='int' (not a real dtype), or passing a Python int like index_dtype=32 instead of a dtype object/string like 'int32'.","commonSituations":"Misreading the API and passing bit-width numbers or dtype names that don't exist; passing the platform default index type from another library.","solutions":["Pass a valid integer dtype string or object: index_dtype=jnp.int32 or 'int32'.","Check the value with dtypes.issubdtype(index_dtype, np.integer) if it comes from config.","Use jnp.argmin/jnp.argmax wrappers, which choose the index dtype automatically."],"exampleFix":"# before\nlax.argmin(x, 0, index_dtype=32)\n# after\nlax.argmin(x, 0, index_dtype=jnp.int32)","handlingStrategy":"type-guard","validationCode":"import numpy as np\nfrom jax import dtypes\nassert dtypes.issubdtype(index_dtype, np.integer), index_dtype\ni = lax.argmin(x, 0, index_dtype=index_dtype)","typeGuard":"def is_integer_dtype(dt):\n    import numpy as np\n    from jax import dtypes\n    return dtypes.issubdtype(dt, np.integer)","tryCatchPattern":null,"preventionTips":["Pass jnp.int32/jnp.int64 objects, not numbers or invented names.","Prefer jnp.argmin/jnp.argmax wrappers unless you need the low-level control."],"tags":["jax","argmin","dtype","type-error"],"backgroundTag":"invalid-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}