jax-ml/jax · error · ValueError

jax.numpy.nanquantile does not support overwrite_input=True

Error message

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

What it means

jnp.nanquantile (and nanpercentile/nanmedian via delegation) rejects overwrite_input=True and out != None for the same immutability reasons as quantile: JAX cannot mutate the input or write into a user buffer.

Source

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

    >>> jnp.quantile(x, q)
    Array([nan, nan, nan], dtype=float32)
    >>> jnp.nanquantile(x, q)
    Array([1.5, 3. , 4.5], dtype=float32)

    Computing weighted quantiles while ignoring NaNs:

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

def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
              method: str, keepdims: bool, squash_nans: bool, weights: Array | None = None) -> Array:
  if method not in ["linear", "lower", "higher", "midpoint", "nearest", "inverted_cdf"]:
    raise ValueError("method can only be 'linear', 'lower', 'higher', 'midpoint', 'nearest' or 'inverted_cdf'")
  if weights is not None:
    if dtypes.issubdtype(weights.dtype, np.complexfloating):
      raise ValueError("Weights cannot be complex types.")
    if method != "inverted_cdf":
      raise NotImplementedError(f"{method} doesn't support weights. Only method 'inverted_cdf' supports weights.")
    a, weights = promote_dtypes_inexact(a, weights)
    if weights.shape != a.shape:
      if axis is None:
        raise ValueError("Weights shape must match 'a' shape when axis is None.")
      ax_tuple = canonicalize_axis_tuple(axis, a.ndim)
      if weights.shape != tuple(a.shape[ax] for ax in ax_tuple):
        raise ValueError(f"Weights shape {weights.shape} must match reduction axes "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Drop overwrite_input and out arguments
Rewrite memory-optimization tricks as explicit sort-based preprocessing under jax.jit

Example fix

// before
np.nanquantile(a, 0.5, overwrite_input=True)
// after
jnp.nanquantile(a, 0.5)

Defensive patterns

Strategy: validation

Validate before calling

kwargs = dict(overwrite_input=False, out=None)
q = jnp.nanquantile(a, q, method=method, **{k: v for k, v in kwargs.items() if v})

Prevention

Remove overwrite_input and out when porting nan-quantile APIs
Whitelist supported kwargs in shared stats wrappers

When it happens

Trigger: Calling jnp.nanquantile(a, q, overwrite_input=True, out=buf), or ported np.nanmedian(x, overwrite_input=True).

Common situations: Porting NaN-aware NumPy statistics code that relied on overwrite_input for large arrays.

Related errors

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