{"record":{"id":"c62121a5c6e301ad","repo":"jax-ml/jax","slug":"expected-fill-values-to-have-dtype-out-dtype-go","errorCode":null,"errorMessage":"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}","messagePattern":"Expected fill values to have dtype (.+?), got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/hijax.py","lineNumber":211,"sourceCode":"      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}\"\n        ) from e\n      if broadcasted != batch_shape:\n        raise ValueError(\n            f\"fill_value shape {fv_aval.shape} cannot be broadcast to \"\n            f\"batch shape {batch_shape} without expanding it.\"\n        )\n    self.in_avals = (a_aval, *fill_value_avals)\n","sourceCodeStart":193,"sourceCodeEnd":229,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/hijax.py#L193-L229","documentation":"Raised by the Nonzero HiJAX primitive when any fill_value's dtype differs from the primitive's integer out_dtype. Fill values pad the returned index arrays, so they must match the index dtype exactly (no implicit casting).","triggerScenarios":"Constructing Nonzero with a float32 fill value aval while out_dtype is int32, e.g. fill_value=jnp.asarray(0.0) with dtype='int32' indices.","commonSituations":"Supplying Python floats or float arrays as fill values; the public nonzero() wrapper converts via jnp.asarray(fv, dtype=out_dtype), so this fires only on direct primitive construction or when dtypes drift.","solutions":["Create fill values with the index dtype: fv = jnp.asarray(0, dtype=out_dtype)","Pass integer scalars (0) rather than floats (0.0)"],"exampleFix":"# before\nfv = jnp.asarray(-1.0)\nprim = Nonzero(aval, core.typeof(fv), size=n, axes=(0,), out_dtype=np.dtype('int32'))\n# after\nfv = jnp.asarray(-1, dtype=jnp.int32)\nprim = Nonzero(aval, core.typeof(fv), size=n, axes=(0,), out_dtype=np.dtype('int32'))","handlingStrategy":"validation","validationCode":"assert all(jnp.asarray(fv).dtype == out_dtype for fv in fill_values)","typeGuard":"def fills_have_dtype(fills, out_dtype) -> bool:\n    return all(jnp.asarray(fv).dtype == out_dtype for fv in fills)","tryCatchPattern":null,"preventionTips":["Create fills with the index dtype: jnp.asarray(-1, dtype=out_dtype)","Use Python ints, not floats, for index padding values"],"tags":["jax","nonzero","fill-value","dtype"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}