What causes "The 'out' argument to jnp.ptp is not supported." in jax-ml/jax?

Remove out and assign the result: rng = jnp.ptp(a, axis=0)

jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.ptp (peak-to-peak) rejects the NumPy-style out= argument because JAX arrays are immutable and cannot be written in place. The out parameter exists only for signature parity.

Source

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

    >>> jnp.ptp(x, axis=1)
    Array([4, 7, 6], dtype=int32)

    To preserve the dimensions of input, you can set ``keepdims=True``.

    >>> jnp.ptp(x, axis=1, keepdims=True)
    Array([[4],
           [7],
           [6]], dtype=int32)
  """
  a = ensure_arraylike("ptp", a)
  return _ptp(a, _ensure_optional_axes(axis), out, keepdims)

@api.jit(static_argnames=('axis', 'keepdims'))
def _ptp(a: Array, axis: Axis = None, out: None = None,
         keepdims: bool = False) -> Array:
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.ptp is not supported.")
  x = amax(a, axis=axis, keepdims=keepdims)
  y = amin(a, axis=axis, keepdims=keepdims)
  return lax.sub(x, y)


@export
@api.jit(static_argnames=('axis', 'keepdims'))
def count_nonzero(a: ArrayLike, axis: Axis = None,
                  keepdims: bool = False) -> Array:
  r"""Return the number of nonzero elements along a given axis.

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

  Args:
    a: input array.
    axis: optional, int or sequence of ints, default=None. Axis along which the
      number of nonzeros are counted. If None, counts within the flattened array.
    keepdims: bool, default=False. If true, reduced axes are left in the result

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove out and assign the result: rng = jnp.ptp(a, axis=0)
Replace buffer-reuse patterns with jit-compiled functions returning the value

Example fix

// before
jnp.ptp(a, out=buf)
// after
buf = jnp.ptp(a)

Defensive patterns

Strategy: validation

Validate before calling

assert out is None or out is ..., 'jnp.ptp does not support out'
result = jnp.ptp(a, axis=axis)

Prevention

Remove out= when porting ptp
Use return-value assignment instead of buffer writes

When it happens

Trigger: Calling jnp.ptp(a, axis=0, out=buf) or passing a third positional argument.

Common situations: Mechanical NumPy→JAX ports that keep out= kwargs; older code targeting buffer reuse on constrained devices.

Related errors

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