jax-ml/jax · error · NotImplementedError

The 'out' argument to jnp.compress is not supported.

Error message

The 'out' argument to jnp.compress is not supported.

What it means

jnp.compress does not support writing results into a preallocated output array, so passing a non-None value for the out parameter raises NotImplementedError. JAX arrays are immutable and functions are pure, so out-style in-place semantics from NumPy are generally unsupported.

Source

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

    Array([[ 1,  3],
           [ 5,  7],
           [ 9, 11]], dtype=int32)

    The optional ``size`` argument lets you specify a static output size so
    that the output is statically-shaped, and so this function can be used
    with transformations like :func:`~jax.jit` and :func:`~jax.vmap`:

    >>> f = lambda c, a: jnp.extract(c, a, size=len(a), fill_value=0)
    >>> mask = (a % 3 == 0)
    >>> jax.vmap(f)(mask, a)
    Array([[ 3,  0,  0,  0],
           [ 6,  0,  0,  0],
           [ 9, 12,  0,  0]], dtype=int32)
  """
  condition_arr, arr, fill_value = util.ensure_arraylike("compress", condition, a, fill_value)
  condition_arr = condition_arr.astype(bool)
  if out is not None:
    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")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove the out argument and use the return value: result = jnp.compress(condition, a)
If buffer reuse was for memory reasons, note JAX's functional model makes it unnecessary
For NumPy interop, drop out= and assign afterwards: buf[:] = jnp.compress(condition, a)

Example fix

// before
np.compress(cond, a, out=buf)
// after
buf = jnp.compress(cond, a)

Defensive patterns

Strategy: type-guard

Validate before calling

assert out is None, 'jnp.compress does not support out='

Try / catch

try:
    result = jnp.compress(cond, a, out=out)
except NotImplementedError:
    result = jnp.compress(cond, a)  # out unsupported

Prevention

Never pass out= to jnp functions; JAX is functional
Grep ported NumPy code for out= arguments
Bind results via assignment, not output buffers

When it happens

Trigger: Calling jnp.compress(condition, a, out=my_array) with any non-None out argument.

Common situations: Porting NumPy code that uses out= to reuse buffers; copy-pasting numpy.compress calls into JAX codebases.

Related errors

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