jax-ml/jax · error · ValueError

fill_value must be a scalar or tuple of scalars; got {fill_v

Error message

fill_value must be a scalar or tuple of scalars; got {fill_value}

What it means

Raised by the public jax.numpy nonzero() function when any fill_value element is not a scalar (after conversion to out_dtype). Fill values pad the fixed-size index outputs and must be 0-dimensional; arrays with any nonzero ndim are rejected.

Source

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

  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)


def einsum(
    subscripts: str,
    /,
    *operands: ArrayLike,
    optimize: str | bool | tuple[tuple[int, ...], ...] = "auto",

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use scalar fill values only: fill_value=-1 or fill_value=(-1, -1)
If a shape-(1,) sneaked in, squeeze it or pass the Python scalar directly
For per-batch behavior, construct the Nonzero primitive with broadcast-compatible fill arrays instead of the public function

Example fix

# before
idx = nonzero(a, size=8, fill_value=jnp.array([-1]))
# after
idx = nonzero(a, size=8, fill_value=-1)

Defensive patterns

Strategy: validation

Validate before calling

fv_tup = fill_value if isinstance(fill_value, tuple) else (fill_value,)
assert all(jnp.asarray(fv).ndim == 0 for fv in fv_tup)

Type guard

def scalar_fills(fill_value) -> bool:
    fv = fill_value if isinstance(fill_value, tuple) else (fill_value,)
    return all(jnp.asarray(v).ndim == 0 for v in fv)

Prevention

Pass Python scalars or 0-d jnp arrays as fill values
Squeeze accidental shape-(1,) arrays: jnp.asarray(x).squeeze()

When it happens

Trigger: nonzero(a, size=8, fill_value=jnp.arange(3)) or fill_value=([0, 1],); any list/array fill_value with ndim > 0.

Common situations: Passing an array of fill values intended per-batch (not supported — fills must be scalars); passing a shape-(1,) array like jnp.array([0.0]) instead of 0.0.

Related errors

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