{"record":{"id":"98cf0ae4b24aadf1","repo":"jax-ml/jax","slug":"out-dtype-must-be-integer-typed-got-out-dtype","errorCode":null,"errorMessage":"out_dtype must be integer typed; got {out_dtype=}","messagePattern":"out_dtype must be integer typed; got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/hijax.py","lineNumber":203,"sourceCode":"  size: int\n  axes: tuple[int, ...]\n  out_dtype: np.dtype\n\n  def __init__(\n      self,\n      a_aval: core.ShapedArray,\n      *fill_value_avals: core.ShapedArray,\n      size: int,\n      axes: tuple[int, ...],\n      out_dtype: np.dtype):\n    if core.is_symbolic_dim(size):\n      pass\n    else:\n      size = operator.index(size)\n      if size < 0:\n        raise ValueError(f\"size must be a positive integer; got {size=}\")\n    if not dtypes.issubdtype(out_dtype, np.integer):\n      raise ValueError(f\"out_dtype must be integer typed; got {out_dtype=}\")\n    if not all(0 <= ax < a_aval.ndim for ax in axes):\n      raise ValueError(f\"axes out of range for array with {a_aval.ndim} dimensions:  {axes=}\")\n    if len(axes) != len(set(axes)):\n      raise ValueError(f\"duplicate axes are not allowed: {axes=}\")\n    if fill_value_avals and len(fill_value_avals) != len(axes):\n      raise ValueError(f\"Expected {len(axes)} fill values, got {len(fill_value_avals)}\")\n    if any(fv.dtype != out_dtype for fv in fill_value_avals):\n      raise ValueError(f\"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}\")\n    batch_shape = tuple(\n        s for i, s in enumerate(a_aval.shape) if i not in axes\n    )\n    for fv_aval in fill_value_avals:\n      try:\n        broadcasted = lax.broadcast_shapes(fv_aval.shape, batch_shape)\n      except ValueError as e:\n        raise ValueError(\n            f\"fill_value shape {fv_aval.shape} is not broadcast-compatible with \"\n            f\"batch shape {batch_shape}\"","sourceCodeStart":185,"sourceCodeEnd":221,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/hijax.py#L185-L221","documentation":"Raised by the Nonzero HiJAX primitive when the output dtype for the returned index arrays is not an integer type. Nonzero returns indices, so out_dtype must be integer-typed (typically int32).","triggerScenarios":"Calling nonzero(a, size=n, dtype='float32') or dtype=np.float64.","commonSituations":"Reusing a dtype config meant for data arrays; assuming dtype refers to the input array's dtype rather than the output indices.","solutions":["Use dtype='int32' or 'int64' (the default is int32)","If float output is needed, cast the returned index arrays after the call"],"exampleFix":"# before\nidx = nonzero(a, size=10, dtype=jnp.float32)\n# after\nidx = nonzero(a, size=10, dtype=jnp.int32)","handlingStrategy":"validation","validationCode":"assert dtypes.issubdtype(np.dtype(dtype), np.integer), dtype","typeGuard":"def is_int_dtype(d) -> bool:\n    return dtypes.issubdtype(np.dtype(d), np.integer)","tryCatchPattern":null,"preventionTips":["Use the default dtype='int32' for nonzero unless indices can exceed 2^31","Don't reuse data-array dtype configs for index outputs"],"tags":["jax","dtype","nonzero","argument-validation"],"backgroundTag":"invalid-dtype-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}