{"record":{"id":"a89b7402870e4070","repo":"jax-ml/jax","slug":"function-dbg-func-src-info-traced-for-dbg-trace-a89b74","errorCode":null,"errorMessage":"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}","messagePattern":"function (.+?) traced for (.+?) returned a mutable array reference of type (.+?)(.+?), but mutable array references cannot be returned\\.(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/partial_eval.py","lineNumber":2203,"sourceCode":"      if list(result_paths) == [\"result\"]: result_paths = [\"\"]  # TODO(mattjj): fix in callee\n      loc = result_paths[i] and f' at output tree path {result_paths[i]}'\n      frame = t._trace.frame\n      v = t.val\n      eqns = frame.get_eqns()\n      # TODO(dougalm): something more efficient\n      eqn = next((e for e in eqns if v in e.outvars), None)\n      if eqn:\n        assert eqn.primitive in (core.ref_p, core.empty_ref_p)\n        origin_info = ('\\n\\nThe returned mutable array was created on line '\n                       f'{source_info_util.summarize(eqn.source_info)}.')\n      elif v in frame.invars:\n        assert isinstance(v, Var)\n        arg_name = dbg.safe_arg_names(len(frame.invars))[frame.invars.index(v)]\n        origin_info = ('\\n\\nThe returned mutable array was passed in as the '\n                       f'argument {arg_name}.')\n      else:\n        origin_info = ''\n      raise ValueError(\n          f\"function {dbg.func_src_info} traced for {dbg.traced_for} returned \"\n          f\"a mutable array reference of type {a.str_short()}{loc}, but \"\n          f\"mutable array references cannot be returned.{origin_info}\")\n\nclass TracerAsName:\n  ref: Any\n  def __init__(self, tracer):\n    self.ref = core.get_referent(tracer)\n  def __eq__(self, other):\n    return isinstance(other, TracerAsName) and self.ref is other.ref\n  def __hash__(self):\n    return id(self.ref)\n\nConst = Any\nVal = Any\n\n\ndef inline_jaxpr_into_trace(","sourceCodeStart":2185,"sourceCodeEnd":2221,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/partial_eval.py#L2185-L2221","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Do not return the Ref; mutate it in place inside the function and return only regular arrays","Return `r[...]` (the read value) instead of `r` if you need the current contents out","Restructure to functional style: pass/return plain arrays and keep state management at the caller"],"exampleFix":"# before\n@jax.jit\ndef f(r, x):\n    r[...] = r[...] + x\n    return r\n\n# after\n@jax.jit\ndef f(r, x):\n    r[...] = r[...] + x\n    return r[...]","handlingStrategy":"validation","validationCode":"from jax.experimental import ref  # or appropriate import\nout = fn(*args)\nassert not any(type(l).__name__ in ('Ref', 'RefValue') for l in jax.tree.leaves(out))","typeGuard":"def returns_no_refs(tree) -> bool:\n    return all(not hasattr(l, 'unsafe_get') for l in jax.tree.leaves(tree))","tryCatchPattern":null,"preventionTips":["Never return Refs; mutate in place and return r[...] or plain arrays","Audit output pytrees when converting stateful code to Ref-based style"],"tags":["jax","mutable-arrays","refs","tracing","return-value"],"backgroundTag":"invalid-return-type-from-traced-function","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}