jax-ml/jax · error · ValueError

condition contains entries that are out of bounds

Error message

condition contains entries that are out of bounds

What it means

When jnp.compress is called without an explicit size, JAX must evaluate the condition concretely and checks that no condition entries lie beyond the length of the (moved) axis. If condition_arr has more entries than arr.shape[0], the excess (extra) is non-empty and ValueError 'condition contains entries that are out of bounds' is raised.

Source

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

    raise NotImplementedError("The 'out' argument to jnp.compress is not supported.")
  if condition_arr.ndim != 1:
    raise ValueError("condition must be a 1D array")
  if axis is None:
    axis = 0
    arr = ravel(arr)
  else:
    arr = moveaxis(arr, axis, 0)
  condition_arr, extra = condition_arr[:arr.shape[0]], condition_arr[arr.shape[0]:]
  arr = arr[:condition_arr.shape[0]]

  if size is None:
    msg = ("The size argument of jnp.compress must be specified in order to use "
           "jnp.compress within JAX transformations like jax.jit, jax.vmap, and "
           "jax.grad. For more information, refer to the jnp.compress documentation.")
    condition_arr = core.concrete_or_error(None, condition_arr, msg)
    extra = core.concrete_or_error(None, extra, msg)
    if extra.any():
      raise ValueError("condition contains entries that are out of bounds")
    result = arr[condition_arr]
  elif not 0 <= size <= arr.shape[0]:
    raise ValueError("size must be positive and not greater than the size of the array axis;"
                     f" got {size=} for a.shape[axis]={arr.shape[0]}")
  else:
    mask = expand_dims(condition_arr, range(1, arr.ndim))
    arr = where(mask, arr, array(fill_value, dtype=arr.dtype))
    result = arr[argsort(condition_arr, stable=True, descending=True)][:size]
  return moveaxis(result, 0, axis)


@export
@api.jit(static_argnames=('rowvar', 'bias', 'ddof', 'dtype'))
def cov(m: ArrayLike, y: ArrayLike | None = None, rowvar: bool = True,
        bias: bool = False, ddof: int | None = None,
        fweights: ArrayLike | None = None,
        aweights: ArrayLike | None = None,
        dtype: DTypeLike | None = None) -> Array:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Trim or correct the condition so its length equals the axis length
2. Pass size explicitly (size=k) to use the padded path which is jit-compatible and bounds-checked differently
3. Check condition.shape[0] == a.shape[axis] before calling

Example fix

// before
jnp.compress(jnp.array([1,0,1,1], bool), jnp.arange(3))  # ValueError
// after
jnp.compress(jnp.array([1,0,1], bool), jnp.arange(3))

cond = jnp.asarray(cond, bool)
assert cond.shape[0] <= a.shape[axis if axis is not None else 0], 'condition longer than axis'
jnp.compress(cond, a, axis=axis)

When it happens

Trigger: jnp.compress(condition, a) where len(condition) > a.shape[axis] (e.g. 5 conditions for an axis of length 3), with size=None so the concrete path is taken.

Common situations: Mismatch between mask length and array length after slicing or filtering data; under jit/vmap the same call instead fails with a concreteness error, a related pitfall.

Related errors

• scan got `length` argument of {} which disagrees with leadin
• reduction axes {axes} contains out-of-bounds indices for {op
• axis argument out of range: {axis=} for {operand.shape=}
• index {i} out of bounds for axis {axis} with size {size} ({n
• index is out of bounds for axis {axis} with size 0

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