{"record":{"id":"2a1c00bc6715609b","repo":"jax-ml/jax","slug":"the-rematted-computation-s-closure-contains-a-muta","errorCode":null,"errorMessage":"the rematted computation's closure contains a mutable array reference of type {v.aval.str_short()} that is not one of the rematted function's inputs, but such refs cannot be saved","messagePattern":"the rematted computation's closure contains a mutable array reference of type (.+?) that is not one of the rematted function's inputs, but such refs cannot be saved","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/ad_checkpoint.py","lineNumber":1059,"sourceCode":"          'the prevent_cse argument to jax.checkpoint'))\n    out_flat = RematTraced(jaxpr, policy, prevent_cse_)(*consts, *args_ft)\n    return out_avals_ft.update(out_flat).unflatten()\n  return decorator\n\ndef _static_argnums(f, argnums, argnames) -> frozenset[int]:\n  argnums = set(argnums)\n  if argnames:\n    sig = api_util.fun_signature(f)\n    assert sig is not None\n    argnums |= set(api_util.infer_argnums_and_argnames(sig, None, argnames)[0])\n  return frozenset(argnums)\n\ndef dce(traced, policy):\n  in_fwd = pe._jaxpr_forwarding(traced.jaxpr)\n  jaxpr = pe.prune_jaxpr_outputs(traced.jaxpr, [f is None for f in in_fwd])\n  for v in jaxpr.outvars:\n    if isinstance(v.aval, AbstractRef):\n      raise ValueError(\n          \"the rematted computation's closure contains a mutable array \"\n          f\"reference of type {v.aval.str_short()} that is not one of the \"\n          \"rematted function's inputs, but such refs cannot be saved\")\n  # dce_jaxpr preserves attached consts (constvars are never pruned).\n  jaxpr, used = pe.dce_jaxpr(jaxpr, True)\n  keep = [u or i in {*in_fwd} for i, u in enumerate(used)]\n  kept_idx = {i: p for p, i in enumerate(i for i, k in enumerate(keep) if k)}\n  in_fwd = tuple(kept_idx[f] if f is not None else None for f in in_fwd)\n  take = tuple(kept_idx[i] for i, u in enumerate(used) if u)\n  keep_res, keep_primals = split_list(keep, [traced._num_consts])\n  res = [r for r, u in zip(traced._consts, keep_res) if u]\n  return keep_primals, Partial(\n      partial(_dced, jaxpr, in_fwd, take, traced.out_tree, policy), res)\n\n@source_info_util.extend_name_stack('rematted_computation')\ndef _dced(jaxpr, in_fwd, take, out_tree, policy, res, *args):\n  ins = [*res, *args]\n  outs = RematTraced(jaxpr, policy)(*[ins[i] for i in take])","sourceCodeStart":1041,"sourceCodeEnd":1077,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ad_checkpoint.py#L1041-L1077","documentation":"PyTreeDef::Compose rejects composing two treedefs whose registries differ. Each PyTreeDef is bound to the PyTreeRegistry it was created under (global or a custom registry), and composition requires them to be identical objects.","triggerScenarios":"outer.Compose(inner) where one treedef came from a custom registry (e.g. forest registry / jax.extend.tree_util) and the other from the global registry, or from two different custom registries.","commonSituations":"Mixing treedefs produced by jax.tree_util (global registry) with treedefs from a custom PyTreeRegistry instance; library code composing treedefs across module boundaries with different registries.","solutions":["Rebuild both treedefs under the same registry instance","Pass the registry explicitly when constructing PyTreeDefs so they match","Register custom node types in the registry the treedefs share"],"exampleFix":"# before\nouter = global_treedef; inner = custom_registry_treedef\nouter.compose(inner)  # raises\n# after\ninner = custom_registry.BuildPyTreeDef(...) rebuilt under outer's registry\nouter.compose(inner)","handlingStrategy":"validation","validationCode":"assert outer.registry() is inner.registry(), 'registry mismatch'\ncomposed = outer.compose(inner)","typeGuard":"def same_registry(a, b) -> bool:\n    return a.registry() is b.registry()","tryCatchPattern":"try:\n    outer.compose(inner)\nexcept ValueError as e:\n    if 'registries' in str(e):\n        inner = _rebuild_under(inner, outer.registry())\n    else:\n        raise","preventionTips":["Thread a single registry object through all treedef construction","Check registry identity before Compose/Tuple","Avoid mixing global-registry and custom-registry treedefs"],"tags":["pytree","jax","registry","compose"],"backgroundTag":"pytree-registry-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}