jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

JAX arrays are immutable, so in-place output buffers like NumPy's out parameter are not supported. jnp.round raises NotImplementedError when out is passed.

Source

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

      nearest integer towards zero.

  Examples:
    >>> x = jnp.array([1.532, 3.267, 6.149])
    >>> jnp.round(x)
    Array([2., 3., 6.], dtype=float32)
    >>> jnp.round(x, decimals=2)
    Array([1.53, 3.27, 6.15], dtype=float32)

    For values exactly halfway between rounded values:

    >>> x1 = jnp.array([10.5, 21.5, 12.5, 31.5])
    >>> jnp.round(x1)
    Array([10., 22., 12., 32.], dtype=float32)
  """
  a = util.ensure_arraylike("round", a)
  decimals = core.concrete_or_error(operator.index, decimals, "'decimals' argument of jnp.round")
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.round is not supported.")
  dtype = a.dtype
  if issubdtype(dtype, np.integer):
    if decimals < 0:
      raise NotImplementedError(
        "integer np.round not implemented for decimals < 0")
    return a  # no-op on integer types

  def _round_float(x: ArrayLike) -> Array:
    if decimals == 0:
      return lax.round(x, lax.RoundingMethod.TO_NEAREST_EVEN)

    # TODO(phawkins): the strategy of rescaling the value isn't necessarily a
    # good one since we may be left with an incorrectly rounded value at the
    # end due to precision problems. As a workaround for float16, convert to
    # float32,
    x = lax.convert_element_type(x, np.float32) if dtype == np.float16 else x
    factor = lax._const(x, 10 ** decimals)
    out = lax.div(lax.round(lax.mul(x, factor),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove out= and use the return value: a = jnp.round(a)
2. If you need buffer semantics, work with a mutable container (e.g. Python list or a numpy array via device_get) and assign back
3. Strip out from kwargs before forwarding to jnp.round

Example fix

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

kwargs.pop('out', None)  # before forwarding to jnp.round

When it happens

Trigger: Calling jnp.round(a, decimals=0, out=buf) — typically code ported from NumPy that reuses a preallocated output array.

Common situations: Translating performance-tuned NumPy code that uses out= to avoid allocations; generic wrapper functions forwarding **kwargs including out.

Related errors

• integer np.round not implemented for decimals < 0
• The 'out' argument to jnp.compress is not supported.
• The 'out' argument to jnp.{name} is not supported.
• The 'out' argument to jnp.logaddexp.reduce is not supported.
• The 'out' argument to jnp.logaddexp2.reduce is not supported

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