jax-ml/jax · error · ValueError

condlist must have length equal to choicelist ({} vs {})

Error message

condlist must have length equal to choicelist ({} vs {})

What it means

jnp.select requires condlist and choicelist to have identical lengths because each condition maps to one choice. A length mismatch means JAX cannot pair conditions with results.

Source

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

    Array([ 10,   2, 300,   0], dtype=int32)

    This is logically equivalent to the following nested ``where`` statement:

    >>> default = 0
    >>> jnp.where(condlist[0],
    ...   choicelist[0],
    ...   jnp.where(condlist[1],
    ...     choicelist[1],
    ...     jnp.where(condlist[2],
    ...       choicelist[2],
    ...       default)))
    Array([ 10,   2, 300,   0], dtype=int32)

    However, for efficiency it is implemented in terms of :func:`jax.lax.select_n`.
  """
  if len(condlist) != len(choicelist):
    msg = "condlist must have length equal to choicelist ({} vs {})"
    raise ValueError(msg.format(len(condlist), len(choicelist)))
  if len(condlist) == 0:
    raise ValueError("condlist must be non-empty")

  util.check_arraylike("select", *condlist, *choicelist, default)
  condlist = [asarray(cond) for cond in condlist]
  choicelist = [asarray(choice) for choice in choicelist]
  default = asarray(default)

  # Put the default at front with condition False because
  # argmax returns zero for an array of False values.
  choicelist = util.promote_dtypes(default, *choicelist)
  conditions = stack(broadcast_arrays(False, *condlist))
  idx = argmax(conditions.astype(bool), axis=0)
  return lax.select_n(*broadcast_arrays(idx, *choicelist))

def is_replicated_or_unreduced(sharding: NamedSharding) -> bool:
  if sharding.spec.partitions:
    return False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure len(condlist) == len(choicelist)
  2. Build both lists together in one loop or zip sources so they stay in sync
  3. Add an assert before calling jnp.select in tests

Example fix

// before
jnp.select([x>0, x<0], [1])
// after
jnp.select([x>0, x<0], [1, -1])
Defensive patterns

Strategy: validation

Validate before calling

assert len(condlist) == len(choicelist), 'condlist/choicelist length mismatch'

Prevention

When it happens

Trigger: Calling jnp.select([c1, c2], [a]) or jnp.select(conds, choices) where the lists were built independently and diverged in length.

Common situations: Building condition and choice lists in separate loops or from dicts with different keys; adding a new condition but forgetting the corresponding branch; version changes adding an extra case.

Related errors


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