jax-ml/jax · error · ValueError

fill_value tuple must have length equal to number of axes ({

Error message

fill_value tuple must have length equal to number of axes ({len(axes)}); got {len(fill_value)}

What it means

Raised by the public jax.numpy nonzero() function when fill_value is given as a tuple whose length differs from the number of axes. One fill value per returned index array (one per axis) is required; a scalar fill_value is broadcast to all axes automatically.

Source

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

  Args:
    a: N-dimensional array.
    size: static integer specifying the number of nonzero entries to return.
    fill_value: optional padding value when ``size`` is specified. Defaults to 0.
    axes: optional tuple of integers specifying the axes to compute the result over.
      Defaults to None (all axes).
    dtype: optional datatype for the returned indices. Defaults to int32.

  Returns:
    Tuple of length ``len(axes)`` containing the indices of each nonzero value.
  """
  a, = core.auto_insert_reshard(a)
  out_dtype = dtypes._maybe_canonicalize_explicit_dtype(np.dtype(dtype), "nonzero")
  axes = util.canonicalize_axis_tuple(axes, np.ndim(a))

  if fill_value is not None:
    if isinstance(fill_value, tuple):
      if len(fill_value) != len(axes):
        raise ValueError(f"fill_value tuple must have length equal to number of axes ({len(axes)}); got {len(fill_value)}")
      fill_value_tup = fill_value
    else:
      fill_value_tup = (fill_value,) * len(axes)
    fill_value_tup = tuple(jnp.asarray(fv, dtype=out_dtype) for fv in fill_value_tup)
    for fv in fill_value_tup:
      if fv.ndim != 0:
        raise ValueError(f"fill_value must be a scalar or tuple of scalars; got {fill_value}")
  else:
    fill_value_tup = ()

  prim = Nonzero(
    core.typeof(a),
    *[core.typeof(fv) for fv in fill_value_tup],
    size=size,
    axes=axes,
    out_dtype=out_dtype,
  )
  return prim(a, *fill_value_tup)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Match lengths: len(fill_value) == len(axes)
Pass a single scalar fill_value (e.g. fill_value=-1) and let it broadcast to every axis
Compute axes explicitly (util.canonicalize_axis_tuple) and derive the tuple length from it

Example fix

# before
idx = nonzero(a, size=8, axes=(0, 1), fill_value=(-1,))
# after
idx = nonzero(a, size=8, axes=(0, 1), fill_value=(-1, -1))
# or simply
idx = nonzero(a, size=8, axes=(0, 1), fill_value=-1)

Defensive patterns

Strategy: validation

Validate before calling

if isinstance(fill_value, tuple):
    assert len(fill_value) == len(axes), (len(fill_value), len(axes))

Type guard

def fill_len_ok(fill_value, n_axes: int) -> bool:
    return fill_value is None or not isinstance(fill_value, tuple) or len(fill_value) == n_axes

Prevention

Pass a scalar fill_value unless per-axis fills are truly needed
When axes is None remember it defaults to all ndim axes — size the tuple accordingly

When it happens

Trigger: nonzero(a, size=8, axes=(0, 1), fill_value=(-1,)) — two axes but one fill value; or fill_value=(0, 0, 0) with the default axes (a.ndim == 2).

Common situations: Adding an axis to the call but forgetting to extend the fill_value tuple; hardcoding a fill_value tuple while axes defaults to all ndim axes of a differently-ranked input.

Related errors

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