jax-ml/jax · error · ValueError

function {dbg.func_src_info} traced for {dbg.traced_for} ret

Error message

function {dbg.func_src_info} traced for {dbg.traced_for} returned a mutable array reference of type {a.str_short()}{loc}, but mutable array references cannot be returned.{origin_info}

What it means

JAX's experimental mutable arrays (Ref/RefValue) can be inputs to traced functions, but cannot be outputs: returning a reference would leak mutable state out of the transform. When the traced result aval is a mutable array reference, JAX raises ValueError, telling you which argument it came from when applicable.

Source

Thrown at jax/_src/interpreters/partial_eval.py:2203

      if list(result_paths) == ["result"]: result_paths = [""]  # TODO(mattjj): fix in callee
      loc = result_paths[i] and f' at output tree path {result_paths[i]}'
      frame = t._trace.frame
      v = t.val
      eqns = frame.get_eqns()
      # TODO(dougalm): something more efficient
      eqn = next((e for e in eqns if v in e.outvars), None)
      if eqn:
        assert eqn.primitive in (core.ref_p, core.empty_ref_p)
        origin_info = ('\n\nThe returned mutable array was created on line '
                       f'{source_info_util.summarize(eqn.source_info)}.')
      elif v in frame.invars:
        assert isinstance(v, Var)
        arg_name = dbg.safe_arg_names(len(frame.invars))[frame.invars.index(v)]
        origin_info = ('\n\nThe returned mutable array was passed in as the '
                       f'argument {arg_name}.')
      else:
        origin_info = ''
      raise ValueError(
          f"function {dbg.func_src_info} traced for {dbg.traced_for} returned "
          f"a mutable array reference of type {a.str_short()}{loc}, but "
          f"mutable array references cannot be returned.{origin_info}")

class TracerAsName:
  ref: Any
  def __init__(self, tracer):
    self.ref = core.get_referent(tracer)
  def __eq__(self, other):
    return isinstance(other, TracerAsName) and self.ref is other.ref
  def __hash__(self):
    return id(self.ref)

Const = Any
Val = Any


def inline_jaxpr_into_trace(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Do not return the Ref; mutate it in place inside the function and return only regular arrays
2. Return `r[...]` (the read value) instead of `r` if you need the current contents out
3. Restructure to functional style: pass/return plain arrays and keep state management at the caller

Example fix

# before
@jax.jit
def f(r, x):
    r[...] = r[...] + x
    return r

# after
@jax.jit
def f(r, x):
    r[...] = r[...] + x
    return r[...]

from jax.experimental import ref  # or appropriate import
out = fn(*args)
assert not any(type(l).__name__ in ('Ref', 'RefValue') for l in jax.tree.leaves(out))

def returns_no_refs(tree) -> bool:
    return all(not hasattr(l, 'unsafe_get') for l in jax.tree.leaves(tree))

When it happens

Trigger: Passing a Ref into a jitted function and returning it (directly or nested in a pytree); `@jax.jit def f(r): ...; return r` or building an output dict containing the Ref.

Common situations: Migrating code from in-place style (flax `ref` methods, `jax.experimental.array_api` mutable state); accidentally including the Ref in an outputs tuple; returning **kwargs that carry the ref.

Related errors

• function {dbg.func_src_info} traced for {dbg.traced_for} ret
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• But the tree structures differ:

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