jax-ml/jax · error · ValueError

fill_value must be a scalar or a tuple of length {arr.ndim};

Error message

fill_value must be a scalar or a tuple of length {arr.ndim}; got {fill_value}

What it means

In jnp.nonzero(..., fill_value=...), each fill value must be a scalar (0-d shape). A fill_value (or element of a tuple fill_value) with non-empty shape is rejected, since output padding entries must be scalars per dimension.

Source

Thrown at jax/_src/numpy/lax_numpy.py:3733

    raise ValueError("Calling nonzero on 0d arrays is not allowed. "
                     "Use jnp.atleast_1d(scalar).nonzero() instead.")
  mask = arr if arr.dtype == bool else (arr != 0)
  calculated_size_ = mask.sum() if size is None else size
  calculated_size: int = core.concrete_dim_or_error(calculated_size_,
    "The size argument of jnp.nonzero must be statically specified "
    "to use jnp.nonzero within JAX transformations.")
  if arr.size == 0 or calculated_size == 0:
    return tuple(array_creation.zeros(calculated_size, int) for dim in arr.shape)
  flat_indices = reductions.cumsum(
      bincount(reductions.cumsum(mask), length=calculated_size))
  strides: np.ndarray = np.cumprod(arr.shape[::-1])[::-1] // arr.shape
  if all(core.is_constant_dim(d) for d in strides):
    strides = strides.astype(flat_indices.dtype)
  out = tuple((flat_indices // stride) % size for stride, size in zip(strides, arr.shape))
  if fill_value is not None:
    fill_value_tup = fill_value if isinstance(fill_value, tuple) else arr.ndim * (fill_value,)
    if any(np.shape(val) != () for val in fill_value_tup):
      raise ValueError(f"fill_value must be a scalar or a tuple of length {arr.ndim}; got {fill_value}")
    fill_mask = arange(calculated_size) >= mask.sum()
    out = tuple(where(fill_mask, fval, entry) for fval, entry in safe_zip(fill_value_tup, out))
  return out


@export
def flatnonzero(a: ArrayLike, *, size: int | None = None,
                fill_value: None | ArrayLike | tuple[ArrayLike, ...] = None) -> Array:
  """Return indices of nonzero elements in a flattened array

  JAX implementation of :func:`numpy.flatnonzero`.

  ``jnp.flatnonzero(x)`` is equivalent to ``nonzero(ravel(a))[0]``. For a full
  discussion of the parameters to this function, refer to :func:`jax.numpy.nonzero`.

  Args:
    a: N-dimensional array.
    size: optional static integer specifying the number of nonzero entries to

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass Python scalars or 0-d values: fill_value=0 or fill_value=(0, 0)
2. Convert arrays: fill_value=int(fv) or fv.item()
3. Verify each element of tuple fill_value has np.shape(val) == ()

Example fix

// before
jnp.nonzero(x, size=5, fill_value=jnp.array([0]))
// after
jnp.nonzero(x, size=5, fill_value=0)

if fill_value is not None:
    if not isinstance(fill_value, tuple):
        fill_value = (fill_value,) * ndim
    fill_value = tuple(v.item() if hasattr(v, 'item') else v for v in fill_value)

def is_scalar_fill(v) -> bool:
    import numpy as np
    return np.shape(v) == ()

When it happens

Trigger: jnp.nonzero(x, size=5, fill_value=jnp.array([0])) — a 1-element array is not a scalar; or a tuple element that is a 1-d array.

Common situations: Passing fill_value loaded from config as an array; reusing index-typed defaults like fill_value=jnp.zeros(1); tuple fill values built from per-dimension arrays.

Related errors

• Expected {len(axes)} fill values, got {len(fill_value_avals)
• Expected fill values to have dtype {out_dtype}, got {fill_va
• fill_value shape {fv_aval.shape} is not broadcast-compatible
• fill_value shape {fv_aval.shape} cannot be broadcast to batc
• fill_value tuple must have length equal to number of axes ({

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