jax-ml/jax · error · ValueError

size and fill_value arguments cannot be used in three-term w

Error message

size and fill_value arguments cannot be used in three-term where function.

What it means

jnp.where has two modes: one-argument (condition only, returning indices via nonzero with size/fill_value) and three-argument (condition, x, y). The size and fill_value kwargs are only meaningful for the one-argument form and are rejected when x and y are provided.

Source

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

    >>> x = jnp.arange(10)
    >>> jnp.where(x > 4)
    (Array([5, 6, 7, 8, 9], dtype=int32),)
    >>> jnp.nonzero(x > 4)
    (Array([5, 6, 7, 8, 9], dtype=int32),)

    When ``x`` and ``y`` are provided, ``where`` selects between them based on
    the specified condition:

    >>> jnp.where(x > 4, x, 0)
    Array([0, 0, 0, 0, 0, 5, 6, 7, 8, 9], dtype=int32)
  """
  if x is None and y is None:
    util.check_arraylike("where", condition)
    return nonzero(condition, size=size, fill_value=fill_value)
  else:
    util.check_arraylike("where", condition, x, y)
    if size is not None or fill_value is not None:
      raise ValueError("size and fill_value arguments cannot be used in "
                       "three-term where function.")
    if x is None or y is None:
      raise ValueError("Either both or neither of the x and y arguments "
                       "should be provided to jax.numpy.where, got "
                       f"{x} and {y}.")
    return util._where(condition, x, y)


@export
def select(
    condlist: Sequence[ArrayLike],
    choicelist: Sequence[ArrayLike],
    default: ArrayLike = 0,
) -> Array:
  """Select values based on a series of conditions.

  JAX implementation of :func:`numpy.select`, implemented in terms
  of :func:`jax.lax.select_n`

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove size and fill_value when using three-argument where
2. If you need a bounded-size index result, use jnp.where(cond) or jnp.nonzero(cond, size=..., fill_value=...) instead

Example fix

// before
jnp.where(cond, x, y, size=5)
// after
jnp.where(cond, x, y)

if size is not None or fill_value is not None:
    assert x is None and y is None, 'size/fill_value only valid for one-arg where'

When it happens

Trigger: Calling jnp.where(cond, x, y, size=10) or jnp.where(cond, x, y, fill_value=-1).

Common situations: Copy-pasting size/fill_value from jitted nonzero-style code into a boolean-select where call; refactoring where(condition) into where(condition, x, y) but leaving the kwargs.

Related errors

• 'trans' value must be 0, 1, or 2, got {trans}
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li
• check_error takes an Error as argument, got type {type(error
• Can't use ``defjvps`` with ``nondiff_argnums``.

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