{"record":{"id":"237647079d4f1fb8","repo":"jax-ml/jax","slug":"condition-contains-entries-that-are-out-of-bounds","errorCode":null,"errorMessage":"condition contains entries that are out of bounds","messagePattern":"condition contains entries that are out of bounds","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":9065,"sourceCode":"    raise NotImplementedError(\"The 'out' argument to jnp.compress is not supported.\")\n  if condition_arr.ndim != 1:\n    raise ValueError(\"condition must be a 1D array\")\n  if axis is None:\n    axis = 0\n    arr = ravel(arr)\n  else:\n    arr = moveaxis(arr, axis, 0)\n  condition_arr, extra = condition_arr[:arr.shape[0]], condition_arr[arr.shape[0]:]\n  arr = arr[:condition_arr.shape[0]]\n\n  if size is None:\n    msg = (\"The size argument of jnp.compress must be specified in order to use \"\n           \"jnp.compress within JAX transformations like jax.jit, jax.vmap, and \"\n           \"jax.grad. For more information, refer to the jnp.compress documentation.\")\n    condition_arr = core.concrete_or_error(None, condition_arr, msg)\n    extra = core.concrete_or_error(None, extra, msg)\n    if extra.any():\n      raise ValueError(\"condition contains entries that are out of bounds\")\n    result = arr[condition_arr]\n  elif not 0 <= size <= arr.shape[0]:\n    raise ValueError(\"size must be positive and not greater than the size of the array axis;\"\n                     f\" got {size=} for a.shape[axis]={arr.shape[0]}\")\n  else:\n    mask = expand_dims(condition_arr, range(1, arr.ndim))\n    arr = where(mask, arr, array(fill_value, dtype=arr.dtype))\n    result = arr[argsort(condition_arr, stable=True, descending=True)][:size]\n  return moveaxis(result, 0, axis)\n\n\n@export\n@api.jit(static_argnames=('rowvar', 'bias', 'ddof', 'dtype'))\ndef cov(m: ArrayLike, y: ArrayLike | None = None, rowvar: bool = True,\n        bias: bool = False, ddof: int | None = None,\n        fweights: ArrayLike | None = None,\n        aweights: ArrayLike | None = None,\n        dtype: DTypeLike | None = None) -> Array:","sourceCodeStart":9047,"sourceCodeEnd":9083,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L9047-L9083","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Trim or correct the condition so its length equals the axis length","Pass size explicitly (size=k) to use the padded path which is jit-compatible and bounds-checked differently","Check condition.shape[0] == a.shape[axis] before calling"],"exampleFix":"// before\njnp.compress(jnp.array([1,0,1,1], bool), jnp.arange(3))  # ValueError\n// after\njnp.compress(jnp.array([1,0,1], bool), jnp.arange(3))","handlingStrategy":"validation","validationCode":"cond = jnp.asarray(cond, bool)\nassert cond.shape[0] <= a.shape[axis if axis is not None else 0], 'condition longer than axis'\njnp.compress(cond, a, axis=axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep condition length <= axis length","Pass size= explicitly under jit/vmap/grad","Slice condition to the axis length: cond[:a.shape[axis]]"],"tags":["jax","out-of-bounds","compress","length-mismatch"],"backgroundTag":"index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}