jax-ml/jax · error · ValueError

with {nc} condition(s), either {nc} or {nc+1} functions are

Error message

with {nc} condition(s), either {nc} or {nc+1} functions are expected; got {nf}

What it means

jnp.piecewise evaluates a piecewise function from a list of conditions and a corresponding list of functions/values. The number of funclist entries must equal the number of conditions or be exactly one more (a default); any other count raises ValueError('with {nc} condition(s), either {nc} or {nc+1} functions are expected; got {nf}').

Source

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

    >>> jnp.piecewise(x, condlist, funclist)
    Array([-3, -2,  -1,  0,  0,  0,  1,  2, 3], dtype=int32)

    ``condlist`` may also be a simple array of scalar conditions, in which case
    the associated function applies to the whole range

    >>> condlist = jnp.array([False, True, False])
    >>> funclist = [lambda x: x * 0, lambda x: x * 10, lambda x: x * 100]
    >>> jnp.piecewise(x, condlist, funclist)
    Array([-40, -30, -20, -10,   0,  10,  20,  30,  40], dtype=int32)
  """
  x_arr = util.ensure_arraylike("piecewise", x)
  nc, nf = len(condlist), len(funclist)
  if nf == nc + 1:
    funclist = funclist[-1:] + funclist[:-1]
  elif nf == nc:
    funclist = [0] + list(funclist)
  else:
    raise ValueError(f"with {nc} condition(s), either {nc} or {nc+1} functions are expected; got {nf}")
  consts = {i: c for i, c in enumerate(funclist) if not callable(c)}
  funcs = {i: f for i, f in enumerate(funclist) if callable(f)}
  return _piecewise(x_arr, asarray(condlist, dtype=bool), consts,
                    frozenset(funcs.items()),  # dict is not hashable.
                    *args, **kw)

@api.jit(static_argnames=['funcs'])
def _piecewise(x: Array, condlist: Array, consts: dict[int, ArrayLike],
               funcs: frozenset[tuple[int, Callable[..., Array]]],
               *args, **kw) -> Array:
  funcdict = dict(funcs)
  funclist = [consts.get(i, funcdict.get(i)) for i in range(len(condlist) + 1)]
  indices = argmax(reductions.cumsum(concatenate(
      [array_creation.zeros_like(condlist[:1]), condlist], 0), 0), 0)
  dtype = x.dtype
  def _call(f):
    return lambda x: f(x, *args, **kw).astype(dtype)
  def _const(v):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make len(funclist) == len(condlist) or len(condlist) + 1
2. Add a default function/value as the extra entry when using nc+1
3. Count both lists before the call

Example fix

// before
jnp.piecewise(x, [x < 0, x >= 0], [lambda x: -x])  # ValueError
// after
jnp.piecewise(x, [x < 0, x >= 0], [lambda x: -x, lambda x: x])

assert len(funclist) in (len(condlist), len(condlist) + 1), \
    f'need {len(condlist)} or {len(condlist)+1} functions, got {len(funclist)}'
jnp.piecewise(x, condlist, funclist)

When it happens

Trigger: jnp.piecewise(x, [c1, c2], [f1]) — 2 conditions with 1 function; or 3 functions for 1 condition.

Common situations: Adding/removing a condition branch without updating funclist; passing a default value without the corresponding extra slot layout piecewise expects.

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• The return value of the policies should be a boolean. Got: {

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