jax-ml/jax · error · ValueError

size must be positive and not greater than the size of the a

Error message

size must be positive and not greater than the size of the array axis; got {size=} for a.shape[axis]={arr.shape[0]}

What it means

When jnp.compress is given an explicit size, that size must satisfy 0 <= size <= arr.shape[axis]. Violating this raises a ValueError showing both the offending size and the axis length.

Source

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

  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:
  r"""Estimate the weighted sample covariance.

  JAX implementation of :func:`numpy.cov`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp or validate size before the call: size = min(size, a.shape[axis])
2. Fix off-by-one in the computation producing size
3. Use size = a.shape[axis] if you meant to include all selected entries

Example fix

// before
jnp.compress(cond, a, size=5)  # a.shape[0] == 3
// after
jnp.compress(cond, a, size=min(5, a.shape[0]))

size = min(max(size, 0), a.shape[axis])
jnp.compress(cond, a, size=size, axis=axis)

When it happens

Trigger: jnp.compress(condition, a, size=n) with n negative or n greater than the length of the selected axis, e.g. size=5 for an axis of length 3.

Common situations: Computing size dynamically (e.g. int(cond.sum()) plus an offset) inside jit where the relationship to axis length isn't enforced; off-by-one errors when size equals axis length + 1.

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/825a1cf3f47febc5. Report an issue: GitHub.