jax-ml/jax · error · ValueError

Expected fill values to have dtype {out_dtype}, got {fill_va

Error message

Expected fill values to have dtype {out_dtype}, got {fill_value_avals}

What it means

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).

Source

Thrown at jax/_src/numpy/hijax.py:211

      size: int,
      axes: tuple[int, ...],
      out_dtype: np.dtype):
    if core.is_symbolic_dim(size):
      pass
    else:
      size = operator.index(size)
      if size < 0:
        raise ValueError(f"size must be a positive integer; got {size=}")
    if not dtypes.issubdtype(out_dtype, np.integer):
      raise ValueError(f"out_dtype must be integer typed; got {out_dtype=}")
    if not all(0 <= ax < a_aval.ndim for ax in axes):
      raise ValueError(f"axes out of range for array with {a_aval.ndim} dimensions:  {axes=}")
    if len(axes) != len(set(axes)):
      raise ValueError(f"duplicate axes are not allowed: {axes=}")
    if fill_value_avals and len(fill_value_avals) != len(axes):
      raise ValueError(f"Expected {len(axes)} fill values, got {len(fill_value_avals)}")
    if any(fv.dtype != out_dtype for fv in fill_value_avals):
      raise ValueError(f"Expected fill values to have dtype {out_dtype}, got {fill_value_avals}")
    batch_shape = tuple(
        s for i, s in enumerate(a_aval.shape) if i not in axes
    )
    for fv_aval in fill_value_avals:
      try:
        broadcasted = lax.broadcast_shapes(fv_aval.shape, batch_shape)
      except ValueError as e:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} is not broadcast-compatible with "
            f"batch shape {batch_shape}"
        ) from e
      if broadcasted != batch_shape:
        raise ValueError(
            f"fill_value shape {fv_aval.shape} cannot be broadcast to "
            f"batch shape {batch_shape} without expanding it."
        )
    self.in_avals = (a_aval, *fill_value_avals)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Create fill values with the index dtype: fv = jnp.asarray(0, dtype=out_dtype)
Pass integer scalars (0) rather than floats (0.0)

Example fix

# before
fv = jnp.asarray(-1.0)
prim = Nonzero(aval, core.typeof(fv), size=n, axes=(0,), out_dtype=np.dtype('int32'))
# after
fv = jnp.asarray(-1, dtype=jnp.int32)
prim = Nonzero(aval, core.typeof(fv), size=n, axes=(0,), out_dtype=np.dtype('int32'))

Defensive patterns

Strategy: validation

Validate before calling

assert all(jnp.asarray(fv).dtype == out_dtype for fv in fill_values)

Type guard

def fills_have_dtype(fills, out_dtype) -> bool:
    return all(jnp.asarray(fv).dtype == out_dtype for fv in fills)

Prevention

Create fills with the index dtype: jnp.asarray(-1, dtype=out_dtype)
Use Python ints, not floats, for index padding values

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c62121a5c6e301ad. Report an issue: GitHub.