jax-ml/jax · error · ValueError

Expected fill values, got

Error message

Expected {len(axes)} fill values, got {len(fill_value_avals)}

What it means

Raised by the Nonzero HiJAX primitive when the number of provided fill_value avals does not equal the number of axes. Each reduced axis needs its own padding value for the returned index arrays when fewer than size nonzeros exist.

Solutions

1. Make len(fill_value) == len(axes), or pass a single scalar fill_value so it is broadcast to all axes
2. Recompute fill values whenever the axes tuple changes

Example fix

# before
prim = Nonzero(aval, fv_aval, size=n, axes=(0, 1), out_dtype=np.int32)
# after
prim = Nonzero(aval, fv_aval, fv_aval, size=n, axes=(0, 1), out_dtype=np.int32)

assert len(fill_value) == len(axes), (len(fill_value), len(axes))

def fill_matches_axes(fill_value, axes) -> bool:
    return fill_value is None or len(fill_value) == len(axes)

When it happens

Trigger: Constructing Nonzero with axes=(0, 1) but supplying only one fill value aval (or three). The public nonzero() wrapper normally catches this earlier with its own message; direct primitive construction hits this one.

Common situations: Mismatch between the axes tuple length and a fill_value list built independently, e.g. after adding an axis but not a fill value.

Related errors

• fill_value must be a scalar or tuple of scalars; got
• fill_value tuple must have length equal to number of axes
• axes out of range for array with
• duplicate axes are not allowed
• Expected fill values to have dtype

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

Appendix: source

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

      a_aval: core.ShapedArray,
      *fill_value_avals: core.ShapedArray,
      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."
        )

View on GitHub (pinned to 1e1c6a8fc0)