jax-ml/jax · error · ValueError

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

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.

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)

Solutions

Make len(fill_value) == len(axes), or pass a single scalar fill_value so it is broadcast to all axes
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)

Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

Prefer a single scalar fill_value so it broadcasts to all axes
Derive fill-value tuples from the axes tuple in the same expression

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

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