jax-ml/jax · error · ValueError

jax.numpy.quantile does not support overwrite_input=True or

Error message

jax.numpy.quantile does not support overwrite_input=True or out != None

What it means

jnp.quantile does not support overwrite_input=True (a NumPy memory optimization that mutates the input during partitioning) or an out= buffer, because JAX arrays are immutable. Either raises ValueError.

Source

Thrown at jax/_src/numpy/reductions.py:2446

    Array([2.25, 4.5 , 6.75], dtype=float32)

    Computing the quartiles using nearest-value interpolation:

    >>> jnp.quantile(x, q, method='nearest')
    Array([2., 4., 7.], dtype=float32)

    Computing weighted quantiles:

    >>> x = jnp.array([1, 2, 3, 4, 5])
    >>> weights = jnp.array([1, 1, 2, 1, 1])
    >>> jnp.quantile(x, 0.5, weights=weights, method='inverted_cdf')
    Array(3., dtype=float32)
  """
  a, q = ensure_arraylike("quantile", a, q)
  if weights is not None:
    weights = ensure_arraylike("quantile", weights)
  if overwrite_input or out is not None:
    raise ValueError("jax.numpy.quantile does not support overwrite_input=True "
                     "or out != None")
  return _quantile(a, q, axis, method, keepdims, False, weights)


@export
@api.jit(static_argnames=('axis', 'overwrite_input', 'keepdims', 'method'))
def nanquantile(a: ArrayLike, q: ArrayLike, axis: int | tuple[int, ...] | None = None,
                out: None = None, overwrite_input: bool = False, method: str = "linear",
                keepdims: bool = False, *, weights: ArrayLike | None = None) -> Array:
  """Compute the quantile of the data along the specified axis, ignoring NaNs.

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

  Args:
    a: N-dimensional array input.
    q: scalar or 1-dimensional array specifying the desired quantiles. ``q``
      should contain floating-point values between ``0.0`` and ``1.0``.
    axis: optional axis or tuple of axes along which to compute the quantile

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Remove overwrite_input and out from the call
  2. If memory matters, operate on jnp.sort/argsort explicitly or rely on jit fusion

Example fix

// before
np.percentile(a, 50, overwrite_input=True)
// after
jnp.percentile(a, 50)
Defensive patterns

Strategy: validation

Validate before calling

def quantile_kwargs(overwrite_input=False, out=None, **kw):
    assert not overwrite_input and out is None, 'jax quantile: out/overwrite_input unsupported'
    return kw
jnp.quantile(a, q, **quantile_kwargs(**user_kwargs))

Prevention

When it happens

Trigger: Calling jnp.quantile(a, q, overwrite_input=True) or jnp.quantile(a, q, out=buf); jnp.median forwards here too, so np.median(x, overwrite_input=True)-style ports also fail.

Common situations: Porting NumPy percentile/median code that used overwrite_input to save memory on large arrays; kwargs passthrough shims.

Related errors


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