Repository: jax-ml/jax
Pull request #27398 (2025-03-25 01:24:54)

[attrs] allow setattr on a previously non-existent attr
Before this change, we handled attrs for initial-style primitives like jit/scan like this:
1. the traceable would form a jaxpr and see what attrs were touched (by jax_getattr or jax_setattr),
2. for each such attr, the traceable would do jax_getattr to get the current value, tree-flatten, pass the flat valuesinto the (pure) bind, get the new values out, tree-unflatten, then jax_setattr the result.

That approach would error if the function called `jax_setattr` to set a previously non-existent attr. That is, this would work at the top level, but not under a `jit`:

```python
from jax.experimental.attrs import jax_setattr
class Thing: ...
thing = Thing()
jax_setattr(thing, 'x', 1.0)
```

This commit makes the same code work under a jit. In short, we use a special value to indicate "does not exist". We just
1. in partial_eval.py's `to_jaxpr`, ensure attrs added during jaxpr formation are deleted, using a special sentinel value `dne_sentinel` to indicate the attribute initially did not exist before tracing;
2. in pjit.py's `_get_states`, when reading initial attr values before the pjit_p bind, if the attribute does not exist we don't try to read it and instead just use `dne_sentinel` as the value, which is a convenient empty pytree;
3. in pjit.py's `_attr_token` for jit caching, when forming the cache key based on the current attr states, we map attrs that don't exist to `dne_sentinel` (rather than just erroring when the attr doesn't exist, as before).

This PR only makes setattr-to-nonexistent-attr work with jit. We'll add scan etc in follow-ups.

(If `jax_getattr` supported the 'default' argument, the code would be a little cleaner since we could avoid the `if hasattr` stuff. And that's probably a useful feature to have anyway. We can add that in a follow-up.)


Patch:
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 07c516fd9..58b97ce2f 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -42,8 +42,8 @@ from jax._src.core import (Trace, Tracer, TraceTag, Jaxpr, Literal, get_aval,
                            mapped_aval, unmapped_aval, DBIdx, InDBIdx, OutDBIdx,
                            InputType, OutputType, get_referent, JaxprEqnContext)
 from jax._src.state.types import AbstractRef, ReadEffect
-from jax._src.tree_util import (PyTreeDef, treedef_tuple,
-                                tree_flatten, tree_structure)
+from jax._src.tree_util import (PyTreeDef, treedef_tuple, tree_flatten,
+                                tree_structure, register_static)
 from jax._src.util import (unzip2, safe_zip, safe_map, toposort, split_list,
                            merge_lists, partition_list, OrderedSet,
                            as_hashable_function, weakref_lru_cache, subs_list,
@@ -1699,7 +1699,7 @@ class JaxprStackFrame:
     jaxpr, constvals = _const_folding_and_forwarding(jaxpr, constvals)
     jaxpr, constvals = _inline_literals(jaxpr, constvals)
     init_trees = [tree_structure(init_val) for init_val in self.attrs_inits]
-    set_states(self.attrs_tracked, self.attrs_inits)
+    set_states(self.attrs_tracked, self.attrs_inits)  # reset to initial values
     return jaxpr, list(constvals), zip(init_trees, end_trees, self.attrs_tracked)
 
   def to_jaxpr2(self, out_tracers: Sequence[core.Tracer],
@@ -2246,11 +2246,15 @@ AttrsTracked = list[tuple[Any, str]]
 AttrStates = list
 def set_states(attrs_tracked: AttrsTracked, vals: AttrStates):
   for ((obj, attr), val) in zip(attrs_tracked, vals):
-    setattr(obj, attr, val)
+    setattr(obj, attr, val) if val is not dne_sentinel else delattr(obj, attr)
 
 def get_states(attrs_tracked: AttrsTracked):
   return [getattr(obj, attr) for (obj, attr) in attrs_tracked]
 
+@register_static
+class DoesNotExist: ...
+dne_sentinel = DoesNotExist()
+
 
 def infer_lambda_input_type(
     axes_specs: Sequence[AbstractedAxesSpec] | None,
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index d690cd6e9..bcdbe6b1b 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -240,10 +240,10 @@ def _set_states(attrs_tracked, vals):
     jax_setattr(obj, attr, val)
 
 def _get_states(attrs_tracked):
-  from jax.experimental.attrs import jax_getattr
+  from jax.experimental.attrs import jax_getattr, dne_sentinel
   vals = []
   for treedef, _, (obj, attr) in attrs_tracked:
-    tree = jax_getattr(obj, attr)
+    tree = jax_getattr(obj, attr) if hasattr(obj, attr) else dne_sentinel
     leaves, treedef_ = tree_flatten(tree)
     assert treedef == treedef_
     vals.extend(leaves)
@@ -1354,11 +1354,11 @@ def _attr_token(
     fun: lu.WrappedFun,
     in_type: core.InputType | tuple[core.AbstractValue, ...]
 ) -> int:
-  from jax.experimental.attrs import jax_getattr
+  from jax.experimental.attrs import jax_getattr, dne_sentinel
   cases = seen_attrs_get(fun, in_type)
   for i, records in enumerate(cases):
     for obj, attr, treedef, avals in records:
-      val = jax_getattr(obj, attr)
+      val = jax_getattr(obj, attr) if hasattr(obj, attr) else dne_sentinel
       vals, treedef_ = tree_flatten(val)
       avals_ = map(core.shaped_abstractify, vals)
       if treedef != treedef_ or avals != avals_: break
@@ -1367,8 +1367,8 @@ def _attr_token(
   return len(cases)
 
 def _attr_update(fun, in_type, i, attrs_tracked):
-  from jax.experimental.attrs import jax_getattr
-  leaves = lambda obj, attr: tree_leaves(jax_getattr(obj, attr))
+  from jax.experimental.attrs import jax_getattr, dne_sentinel
+  leaves = lambda obj, attr: tree_leaves(jax_getattr(obj, attr) if hasattr(obj, attr) else dne_sentinel)
   records = [(obj, attr, init_tree, map(core.shaped_abstractify, leaves(obj, attr)))
              for init_tree, _, (obj, attr) in attrs_tracked]
   cases = seen_attrs_get(fun, in_type)
diff --git a/jax/experimental/attrs.py b/jax/experimental/attrs.py
index 4e1dc4b8f..bb4c7bf83 100644
--- a/jax/experimental/attrs.py
+++ b/jax/experimental/attrs.py
@@ -36,6 +36,7 @@ JaxVal = Any
 Pytree = Any
 
 register = api_util.register_class_with_attrs
+dne_sentinel = pe.dne_sentinel
 
 def jax_getattr(obj: Any, attr: str):
   with core.take_current_trace() as t:
@@ -65,7 +66,7 @@ def _ensure_tracked(trace: pe.DynamicJaxprTrace, obj: Any, attr: str):
     return tracer
 
   if (obj, attr) not in frame.attrs_tracked:
-    init_val = getattr(obj, attr)
+    init_val = getattr(obj, attr, dne_sentinel)
     frame.attrs_inits.append(init_val)
     init_vals, init_tree = tree_flatten(init_val)
     tracers = map(new_tracer, init_vals)
diff --git a/jaxlib/xla/BUILD b/jaxlib/xla/BUILD
index 2edc183bc..592d9d1c2 100644
--- a/jaxlib/xla/BUILD
+++ b/jaxlib/xla/BUILD
@@ -14,16 +14,13 @@
 
 load(
     "//jaxlib:jax.bzl",
-    "cc_proto_library",
     "if_oss",
-    "jax_visibility",
     "nanobind_extension",
     "py_deps",
     "py_strict_library",
     "py_strict_test",
     "pytype_strict_library",
 )
-# Placeholder: load proto_library
 
 licenses(["notice"])
 
@@ -53,7 +50,6 @@ nanobind_extension(
         ":mlir",
         ":pjit",
         ":pmap_lib",
-        ":pytree",
         ":sdy",
         ":weakref_lru_cache",
         ":xla_compiler",
@@ -107,6 +103,7 @@ nanobind_extension(
         "@xla//xla/python:profiler",
         "@xla//xla/python:py_client",
         "@xla//xla/python:python_ref_manager",
+        "@xla//xla/python:pytree",
         "@xla//xla/python:refine_polymorphic_shapes",
         "@xla//xla/python:traceback",
         "@xla//xla/python:types",
@@ -253,7 +250,6 @@ cc_library(
     ],
     features = ["-use_header_modules"],
     deps = [
-        ":pytree",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:inlined_vector",
@@ -274,6 +270,7 @@ cc_library(
         "@xla//xla/python:nb_helpers",
         "@xla//xla/python:py_client",
         "@xla//xla/python:python_ref_manager",
+        "@xla//xla/python:pytree",
         "@xla//xla/python:types",
         "@xla//xla/tsl/platform:logging",
     ],
@@ -331,7 +328,6 @@ cc_library(
     deps = [
         ":config",
         ":jax_jit",
-        ":pytree",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -356,6 +352,7 @@ cc_library(
         "@xla//xla/python:nb_numpy",
         "@xla//xla/python:py_client",
         "@xla//xla/python:python_ref_manager",
+        "@xla//xla/python:pytree",
         "@xla//xla/python:traceback",
         "@xla//xla/python/ifrt",
         "@xla//xla/tsl/concurrency:ref_count",
@@ -379,7 +376,6 @@ cc_library(
     deps = [
         ":config",
         ":jax_jit",
-        ":pytree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/hash",
@@ -403,6 +399,7 @@ cc_library(
         "@xla//xla/python:nb_numpy",
         "@xla//xla/python:py_client",
         "@xla//xla/python:python_ref_manager",
+        "@xla//xla/python:pytree",
         "@xla//xla/python:traceback",
         "@xla//xla/python:types",
         "@xla//xla/python/ifrt",
@@ -414,44 +411,6 @@ cc_library(
     ],
 )
 
-proto_library(
-    name = "pytree_proto",
-    srcs = ["pytree.proto"],
-)
-
-cc_proto_library(
-    name = "pytree_cc_proto",
-    deps = [":pytree_proto"],
-)
-
-cc_library(
-    name = "pytree",
-    srcs = ["pytree.cc"],
-    hdrs = ["pytree.h"],
-    compatible_with = [],
-    copts = [
-        "-fexceptions",
-        "-fno-strict-aliasing",
-    ],
-    features = ["-use_header_modules"],
-    visibility = jax_visibility("jaxlib/xla/pytree"),
-    deps = [
-        ":pytree_cc_proto",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:inlined_vector",
-        "@com_google_absl//absl/hash",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/types:span",
-        "@nanobind",
-        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
-        "@xla//xla/pjrt:exceptions",
-        "@xla//xla/python:nb_class_ptr",
-        "@xla//xla/tsl/platform:logging",
-    ],
-)
-
 cc_library(
     name = "sdy",
     srcs = ["sdy.cc"],
diff --git a/jaxlib/xla/jax_jit.cc b/jaxlib/xla/jax_jit.cc
index 23abe9a84..754272a07 100644
--- a/jaxlib/xla/jax_jit.cc
+++ b/jaxlib/xla/jax_jit.cc
@@ -53,13 +53,13 @@ limitations under the License.
 #include "nanobind/stl/string.h"  // IWYU pragma: keep
 #include "nanobind/stl/string_view.h"  // IWYU pragma: keep
 #include "nanobind/stl/vector.h"  // IWYU pragma: keep
-#include "jaxlib/xla/pytree.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/python/nb_absl_inlined_vector.h"  // IWYU pragma: keep
 #include "xla/python/nb_absl_span.h"  // IWYU pragma: keep
 #include "xla/python/py_values.h"
+#include "xla/python/pytree.h"
 #include "xla/python/sharding.h"
 #include "xla/python/types.h"
 #include "xla/tsl/platform/logging.h"
diff --git a/jaxlib/xla/jax_jit.h b/jaxlib/xla/jax_jit.h
index a000ef677..303d7e694 100644
--- a/jaxlib/xla/jax_jit.h
+++ b/jaxlib/xla/jax_jit.h
@@ -34,11 +34,11 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "nanobind/nanobind.h"
-#include "jaxlib/xla/pytree.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/python/nb_helpers.h"
 #include "xla/python/py_values.h"
 #include "xla/python/python_ref_manager.h"
+#include "xla/python/pytree.h"
 #include "xla/python/sharding.h"
 #include "xla/tsl/platform/logging.h"
 
diff --git a/jaxlib/xla/pjit.cc b/jaxlib/xla/pjit.cc
index a13d3f0b5..96056708c 100644
--- a/jaxlib/xla/pjit.cc
+++ b/jaxlib/xla/pjit.cc
@@ -52,7 +52,6 @@ limitations under the License.
 #include "nanobind/stl/vector.h"  // IWYU pragma: keep
 #include "jaxlib/xla/config.h"
 #include "jaxlib/xla/jax_jit.h"
-#include "jaxlib/xla/pytree.h"
 #include "xla/layout.h"
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/lru_cache.h"
@@ -70,6 +69,7 @@ limitations under the License.
 #include "xla/python/py_executable.h"
 #include "xla/python/py_values.h"
 #include "xla/python/python_ref_manager.h"
+#include "xla/python/pytree.h"
 #include "xla/python/sharding.h"
 #include "xla/python/traceback.h"
 #include "xla/tsl/concurrency/ref_count.h"
diff --git a/jaxlib/xla/pmap_lib.cc b/jaxlib/xla/pmap_lib.cc
index c6849f8c2..5582eccf4 100644
--- a/jaxlib/xla/pmap_lib.cc
+++ b/jaxlib/xla/pmap_lib.cc
@@ -46,7 +46,6 @@ limitations under the License.
 #include "nanobind/stl/vector.h"  // IWYU pragma: keep
 #include "jaxlib/xla/config.h"
 #include "jaxlib/xla/jax_jit.h"
-#include "jaxlib/xla/pytree.h"
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/python/ifrt/array.h"
@@ -65,6 +64,7 @@ limitations under the License.
 #include "xla/python/py_executable.h"
 #include "xla/python/py_values.h"
 #include "xla/python/python_ref_manager.h"
+#include "xla/python/pytree.h"
 #include "xla/python/sharded_device_array.h"
 #include "xla/python/sharding.h"
 #include "xla/python/to_ifrt_sharding.h"
diff --git a/jaxlib/xla/pytree.cc b/jaxlib/xla/pytree.cc
deleted file mode 100644
index dd5a0bd9c..000000000
--- a/jaxlib/xla/pytree.cc
+++ /dev/null
@@ -1,1825 +0,0 @@
-/* Copyright 2019 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-// Caution: this code uses exceptions. The exception use is local to the
-// binding code and the idiomatic way to emit Python exceptions.
-
-#include "jaxlib/xla/pytree.h"
-
-#include <Python.h>
-
-#include <algorithm>
-#include <cstddef>
-#include <cstdint>
-#include <iterator>
-#include <limits>
-#include <memory>
-#include <optional>
-#include <stdexcept>
-#include <string>
-#include <utility>
-#include <vector>
-
-#include "absl/algorithm/container.h"
-#include "absl/container/flat_hash_map.h"
-#include "absl/container/inlined_vector.h"
-#include "absl/hash/hash.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_format.h"
-#include "absl/strings/str_join.h"
-#include "absl/strings/string_view.h"
-#include "absl/types/span.h"
-#include "nanobind/nanobind.h"
-#include "nanobind/stl/optional.h"  // IWYU pragma: keep
-#include "nanobind/stl/pair.h"  // IWYU pragma: keep
-#include "nanobind/stl/shared_ptr.h"  // IWYU pragma: keep
-#include "nanobind/stl/string.h"  // IWYU pragma: keep
-#include "nanobind/stl/string_view.h"  // IWYU pragma: keep
-#include "nanobind/stl/tuple.h"  // IWYU pragma: keep
-#include "nanobind/stl/vector.h"  // IWYU pragma: keep
-#include "jaxlib/xla/pytree.pb.h"
-#include "xla/pjrt/exceptions.h"
-#include "xla/python/nb_class_ptr.h"
-#include "xla/tsl/platform/logging.h"
-
-namespace xla {
-
-namespace nb = nanobind;
-
-constexpr int kSequenceKeyHashSalt = 1;
-constexpr int kFlattenedIndexKeyHashSalt = 42;
-
-PyTreeRegistry::PyTreeRegistry(bool enable_none, bool enable_tuple,
-                               bool enable_namedtuple, bool enable_list,
-                               bool enable_dict) {
-  auto add_builtin_type = [&](PyTypeObject* type_obj, PyTreeKind kind) {
-    nb::object type =
-        nb::borrow<nb::object>(reinterpret_cast<PyObject*>(type_obj));
-    auto registration = std::make_unique<Registration>();
-    registration->kind = kind;
-    registration->type = type;
-    CHECK(registrations_.emplace(type, std::move(registration)).second);
-  };
-  if (enable_none) {
-    add_builtin_type(Py_TYPE(Py_None), PyTreeKind::kNone);
-  }
-  if (enable_tuple) {
-    add_builtin_type(&PyTuple_Type, PyTreeKind::kTuple);
-  }
-  enable_namedtuple_ = enable_namedtuple;
-  if (enable_list) {
-    add_builtin_type(&PyList_Type, PyTreeKind::kList);
-  }
-  if (enable_dict) {
-    add_builtin_type(&PyDict_Type, PyTreeKind::kDict);
-  }
-}
-
-void PyTreeRegistry::Register(
-    nb::object type, nb::callable to_iterable, nb::callable from_iterable,
-    std::optional<nb::callable> to_iterable_with_keys) {
-  auto registration = std::make_unique<Registration>();
-  registration->kind = PyTreeKind::kCustom;
-  registration->type = type;
-  registration->to_iterable = std::move(to_iterable);
-  registration->from_iterable = std::move(from_iterable);
-  registration->to_iterable_with_keys = std::move(to_iterable_with_keys);
-  nb::ft_lock_guard lock(mu_);
-  auto it = registrations_.emplace(type, std::move(registration));
-  if (!it.second) {
-    throw std::invalid_argument(
-        absl::StrFormat("Duplicate custom PyTreeDef type registration for %s.",
-                        nb::cast<absl::string_view>(nb::repr(type))));
-  }
-}
-
-void PyTreeRegistry::RegisterDataclass(nb::object type,
-                                       std::vector<nb::str> data_fields,
-                                       std::vector<nb::str> meta_fields) {
-  auto registration = std::make_unique<Registration>();
-  registration->kind = PyTreeKind::kDataclass;
-  registration->type = type;
-  registration->data_fields = std::move(data_fields);
-  registration->meta_fields = std::move(meta_fields);
-  nb::ft_lock_guard lock(mu_);
-  auto it = registrations_.emplace(type, std::move(registration));
-  if (!it.second) {
-    throw std::invalid_argument(absl::StrFormat(
-        "Duplicate custom dataclass PyTreeDef type registration for %s.",
-        nb::cast<absl::string_view>(nb::repr(std::move(type)))));
-  }
-}
-
-std::pair<nanobind::iterable, nanobind::object>
-PyTreeRegistry::Registration::ToIterable(nanobind::handle o) const {
-  nb::object out = to_iterable(o);
-  nb::tuple leaves_and_aux_data;
-  if (!nb::try_cast<nb::tuple>(out, leaves_and_aux_data) ||
-      leaves_and_aux_data.size() != 2) {
-    throw std::invalid_argument(absl::StrCat(
-        "The to_iterable function for a custom PyTree node should return "
-        "a (children, aux_data) tuple, got ",
-        nb::cast<absl::string_view>(nb::repr(out))));
-  }
-  nb::iterable leaves;
-  if (!nb::try_cast<nb::iterable>(leaves_and_aux_data[0], leaves)) {
-    throw std::invalid_argument(absl::StrCat(
-        "The to_iterable function for a custom PyTree node should return "
-        "a (children, aux_data) tuple where 'children' is iterable, "
-        "got ",
-        nb::cast<absl::string_view>(nb::repr(out))));
-  }
-  return std::make_pair(std::move(leaves), nb::object(leaves_and_aux_data[1]));
-}
-
-std::pair<std::vector<std::pair<nb::object, nb::object>>, nb::object>
-PyTreeRegistry::Registration::ToIterableWithKeys(nb::handle o) const {
-  // Backwards compatibility case: return dummy FlattenedIndexKey for each leaf.
-  std::vector<std::pair<nb::object, nb::object>> result;
-  if (!to_iterable_with_keys.has_value()) {
-    auto [leaves, aux_data] = ToIterable(o);
-    for (nb::handle leaf : leaves) {
-      result.push_back(std::make_pair(
-          make_nb_class<FlattenedIndexKey>(result.size()), nb::borrow(leaf)));
-    }
-    return std::make_pair(std::move(result), std::move(aux_data));
-  }
-  nb::object out = to_iterable_with_keys.value()(o);
-  nb::tuple leaves_and_aux_data;
-  if (!nb::try_cast<nb::tuple>(out, leaves_and_aux_data) ||
-      leaves_and_aux_data.size() != 2) {
-    throw std::invalid_argument(absl::StrCat(
-        "The to_iterable_with_keys function for a custom PyTree "
-        "node should return a (key_leaf_pairs, aux_data) tuple, got ",
-        nb::cast<absl::string_view>(nb::repr(out))));
-  }
-  nb::iterable key_leaf_pairs;
-  if (!nb::try_cast<nb::iterable>(leaves_and_aux_data[0], key_leaf_pairs)) {
-    throw std::invalid_argument(absl::StrCat(
-        "The to_iterable_with_keys function for a custom PyTree node should "
-        "return a (key_leaf_pairs, aux_data) tuple where 'key_leaf_pairs' is "
-        "iterable, got ",
-        nb::cast<absl::string_view>(nb::repr(leaves_and_aux_data))));
-  }
-  for (nb::handle key_leaf_pair : key_leaf_pairs) {
-    nb::tuple key_leaf_pair_tuple;
-    if (!nb::try_cast<nb::tuple>(key_leaf_pair, key_leaf_pair_tuple) ||
-        key_leaf_pair_tuple.size() != 2) {
-      throw std::invalid_argument(absl::StrCat(
-          "The to_iterable_with_keys function for a custom PyTree node should "
-          "return a (key_leaf_pairs, aux_data) tuple where 'child",
-          nb::cast<absl::string_view>(nb::repr(key_leaf_pair))));
-    }
-    result.push_back(std::make_pair(nb::borrow(key_leaf_pair_tuple[0]),
-                                    nb::borrow(key_leaf_pair_tuple[1])));
-  }
-  return std::make_pair(std::move(result), nb::object(leaves_and_aux_data[1]));
-}
-
-int PyTreeRegistry::Registration::tp_traverse(visitproc visit, void* arg) {
-  Py_VISIT(type.ptr());
-  Py_VISIT(to_iterable.ptr());
-  Py_VISIT(from_iterable.ptr());
-  for (const auto& field : data_fields) {
-    Py_VISIT(field.ptr());
-  }
-  for (const auto& field : meta_fields) {
-    Py_VISIT(field.ptr());
-  }
-  return 0;
-}
-
-// Computes the node kind of a given Python object.
-PyTreeKind PyTreeRegistry::KindOfObject(
-    nb::handle obj, PyTreeRegistry::Registration const** custom) const {
-  const PyTreeRegistry::Registration* registration = Lookup(obj.type());
-  if (registration) {
-    if (registration->kind == PyTreeKind::kCustom ||
-        registration->kind == PyTreeKind::kDataclass) {
-      *custom = registration;
-    } else {
-      *custom = nullptr;
-    }
-    return registration->kind;
-  } else if (nb::isinstance<nb::tuple>(obj) && nb::hasattr(obj, "_fields")) {
-    // We can only identify namedtuples heuristically, here by the presence of
-    // a _fields attribute.
-    return PyTreeKind::kNamedTuple;
-  } else {
-    return PyTreeKind::kLeaf;
-  }
-}
-
-/*static*/ const PyTreeRegistry::Registration* PyTreeRegistry::Lookup(
-    nb::handle type) const {
-  nb::ft_lock_guard lock(mu_);
-  auto it = registrations_.find(type);
-  return it == registrations_.end() ? nullptr : it->second.get();
-}
-
-/*static*/ std::vector<nb::object> GetSortedPyDictKeys(PyObject* py_dict) {
-  std::vector<nb::object> keys;
-  keys.reserve(PyDict_Size(py_dict));
-  PyObject* key;
-  Py_ssize_t pos = 0;
-  while (PyDict_Next(py_dict, &pos, &key, /*value=*/nullptr)) {
-    keys.push_back(nb::borrow<nb::object>(key));
-  }
-
-  try {
-    std::stable_sort(
-        keys.begin(), keys.end(), [](const nb::object& a, const nb::object& b) {
-          int cmp = PyObject_RichCompareBool(a.ptr(), b.ptr(), Py_LT);
-          if (cmp == -1) {
-            throw nb::python_error();
-          }
-          return cmp;
-        });
-  } catch (nb::python_error& e) {
-    nb::raise_from(e, PyExc_ValueError,
-                   "Comparator raised exception while sorting pytree "
-                   "dictionary keys.");
-  }
-  return keys;
-}
-
-/*static*/ bool IsSortedPyDictKeysEqual(absl::Span<const nb::object> lhs,
-                                        absl::Span<const nb::object> rhs) {
-  if (lhs.size() != rhs.size()) {
-    return false;
-  }
-  for (int i = 0; i < lhs.size(); ++i) {
-    if (lhs[i].not_equal(rhs[i])) {
-      return false;
-    }
-  }
-  return true;
-}
-
-bool PyTreeDef::operator==(const PyTreeDef& other) const {
-  if (traversal_.size() != other.traversal_.size()) {
-    return false;
-  }
-  for (size_t i = 0; i < traversal_.size(); ++i) {
-    const Node& a = traversal_[i];
-    const Node& b = other.traversal_[i];
-    if (a.kind != b.kind || a.arity != b.arity ||
-        (a.node_data.ptr() == nullptr) != (b.node_data.ptr() == nullptr) ||
-        (a.sorted_dict_keys.size() != b.sorted_dict_keys.size()) ||
-        a.custom != b.custom) {
-      return false;
-    }
-    if (a.node_data && a.node_data.not_equal(b.node_data)) {
-      return false;
-    }
-    if (!IsSortedPyDictKeysEqual(a.sorted_dict_keys, b.sorted_dict_keys)) {
-      return false;
-    }
-    // We don't need to test equality of num_leaves and num_nodes since they
-    // are derivable from the other node data.
-  }
-  return true;
-}
-
-nb::object PyTreeRegistry::FlattenOneLevel(nb::handle x) const {
-  return FlattenOneLevelImpl(x, /*with_keys=*/false);
-}
-
-nb::object PyTreeRegistry::FlattenOneLevelWithKeys(nb::handle x) const {
-  return FlattenOneLevelImpl(x, /*with_keys=*/true);
-}
-
-nb::object PyTreeRegistry::FlattenOneLevelImpl(nb::handle x,
-                                               bool with_keys) const {
-  PyTreeRegistry::Registration const* custom;
-  PyTreeKind kind = KindOfObject(x, &custom);
-  switch (kind) {
-    case PyTreeKind::kNone:
-      return nb::make_tuple(nb::make_tuple(), nb::none());
-    case PyTreeKind::kTuple: {
-      if (with_keys) {
-        auto size = PyTuple_GET_SIZE(x.ptr());
-        nb::object key_leaves = nb::steal(PyTuple_New(size));
-        for (int i = 0; i < size; ++i) {
-          nb::object key = make_nb_class<SequenceKey>(i);
-          nb::object value =
-              nb::borrow<nb::object>(PyTuple_GET_ITEM(x.ptr(), i));
-          PyTuple_SET_ITEM(key_leaves.ptr(), i,
-                           nb::make_tuple(key, value).release().ptr());
-        }
-        return nb::make_tuple(std::move(key_leaves), nb::none());
-      }
-      return nb::make_tuple(nb::borrow(x), nb::none());
-    }
-    case PyTreeKind::kList: {
-      if (with_keys) {
-        auto size = PyList_GET_SIZE(x.ptr());
-        nb::object key_leaves = nb::steal(PyTuple_New(size));
-        for (int i = 0; i < size; ++i) {
-          nb::object key = make_nb_class<SequenceKey>(i);
-          nb::object value =
-              nb::borrow<nb::object>(PyList_GET_ITEM(x.ptr(), i));
-          PyTuple_SET_ITEM(key_leaves.ptr(), i,
-                           nb::make_tuple(key, value).release().ptr());
-        }
-        return nb::make_tuple(std::move(key_leaves), nb::none());
-      }
-      return nb::make_tuple(nb::borrow(x), nb::none());
-    }
-    case PyTreeKind::kDict: {
-      nb::dict dict = nb::borrow<nb::dict>(x);
-      std::vector<nb::object> sorted_keys = GetSortedPyDictKeys(dict.ptr());
-      nb::tuple keys = nb::steal<nb::tuple>(PyTuple_New(sorted_keys.size()));
-      nb::tuple values = nb::steal<nb::tuple>(PyTuple_New(sorted_keys.size()));
-      for (size_t i = 0; i < sorted_keys.size(); ++i) {
-        nb::object& key = sorted_keys[i];
-        nb::object value = nb::object(dict[key]);
-        if (with_keys) {
-          value = nb::make_tuple(make_nb_class<DictKey>(key), value);
-        }
-        PyTuple_SET_ITEM(values.ptr(), i, value.release().ptr());
-        PyTuple_SET_ITEM(keys.ptr(), i, sorted_keys[i].release().ptr());
-      }
-      return nb::make_tuple(std::move(values), std::move(keys));
-    }
-    case PyTreeKind::kNamedTuple: {
-      nb::tuple in = nb::borrow<nb::tuple>(x);
-      nb::list out;
-      if (with_keys) {
-        // Get key names from NamedTuple fields.
-        nb::tuple fields;
-        if (!nb::try_cast<nb::tuple>(nb::getattr(in, "_fields"), fields) ||
-            in.size() != fields.size()) {
-          throw std::invalid_argument(
-              "A namedtuple's _fields attribute should have the same size as "
-              "the tuple.");
-        }
-        auto field_iter = fields.begin();
-        for (nb::handle entry : in) {
-          out.append(nb::make_tuple(
-              make_nb_class<GetAttrKey>(nb::str(*field_iter)), entry));
-        }
-        return nb::make_tuple(std::move(out), x.type());
-      }
-      for (size_t i = 0; i < in.size(); ++i) {
-        out.append(in[i]);
-      }
-      return nb::make_tuple(std::move(out), x.type());
-    }
-    case PyTreeKind::kCustom: {
-      if (with_keys) {
-        auto [leaves, aux_data] = custom->ToIterableWithKeys(x);
-        return nb::make_tuple(std::move(leaves), std::move(aux_data));
-      }
-      auto [leaves, aux_data] = custom->ToIterable(x);
-      return nb::make_tuple(std::move(leaves), std::move(aux_data));
-    }
-    case PyTreeKind::kDataclass: {
-      auto data_size = custom->data_fields.size();
-      nb::list leaves = nb::steal<nb::list>(PyList_New(data_size));
-      for (int leaf = 0; leaf < data_size; ++leaf) {
-        nb::object value = nb::getattr(x, custom->data_fields[leaf]);
-        if (with_keys) {
-          value = nb::make_tuple(
-              make_nb_class<GetAttrKey>(custom->data_fields[leaf]), value);
-        }
-        PyList_SET_ITEM(leaves.ptr(), leaf, value.release().ptr());
-      }
-      auto meta_size = custom->meta_fields.size();
-      nb::object aux_data = nb::steal(PyTuple_New(meta_size));
-      for (int meta_leaf = 0; meta_leaf < meta_size; ++meta_leaf) {
-        PyTuple_SET_ITEM(
-            aux_data.ptr(), meta_leaf,
-            nb::getattr(x, custom->meta_fields[meta_leaf]).release().ptr());
-      }
-      return nb::make_tuple(std::move(leaves), std::move(aux_data));
-    }
-    default:
-      DCHECK(kind == PyTreeKind::kLeaf);
-      return nb::none();
-  }
-}
-
-/* static */ PyType_Slot PyTreeRegistry::slots_[] = {
-    {Py_tp_traverse, (void*)PyTreeRegistry::tp_traverse},
-    {Py_tp_clear, (void*)PyTreeRegistry::tp_clear},
-    {0, nullptr},
-};
-
-/* static */ int PyTreeRegistry::tp_traverse(PyObject* self, visitproc visit,
-                                             void* arg) {
-  PyTreeRegistry* registry = nb::inst_ptr<PyTreeRegistry>(self);
-  Py_VISIT(Py_TYPE(self));
-  nb::ft_lock_guard lock(registry->mu_);
-  for (const auto& [key, value] : registry->registrations_) {
-    Py_VISIT(key.ptr());
-    int rval = value->tp_traverse(visit, arg);
-    if (rval != 0) {
-      return rval;
-    }
-  }
-  return 0;
-}
-
-/* static */ int PyTreeRegistry::tp_clear(PyObject* self) {
-  PyTreeRegistry* registry = nb::inst_ptr<PyTreeRegistry>(self);
-  nb::ft_lock_guard lock(registry->mu_);
-  registry->registrations_.clear();
-  return 0;
-}
-
-/* static */ PyType_Slot DictKey::slots_[] = {
-    {Py_tp_traverse, (void*)DictKey::tp_traverse},
-    {Py_tp_clear, (void*)DictKey::tp_clear},
-    {0, nullptr},
-};
-
-/* static */ int DictKey::tp_traverse(PyObject* self, visitproc visit,
-                                      void* arg) {
-  DictKey* key = nb::inst_ptr<DictKey>(self);
-  Py_VISIT(key->key_.ptr());
-  return 0;
-}
-
-/* static */ int DictKey::tp_clear(PyObject* self) {
-  DictKey* dictkey = nb::inst_ptr<DictKey>(self);
-  nb::object tmp;
-  std::swap(tmp, dictkey->key_);
-  return 0;
-}
-
-std::string SequenceKey::ToString() const {
-  return absl::StrFormat("[%d]", idx_);
-}
-
-std::string SequenceKey::ToReprString() const {
-  return absl::StrFormat("SequenceKey(idx=%d)", idx_);
-}
-
-std::string DictKey::ToString() const {
-  return absl::StrFormat("[%s]", nb::cast<absl::string_view>(nb::repr(key_)));
-}
-
-std::string DictKey::ToReprString() const {
-  return absl::StrFormat("DictKey(key=%s)",
-                         nb::cast<absl::string_view>(nb::repr(key_)));
-}
-
-std::string GetAttrKey::ToString() const {
-  return absl::StrFormat(".%s", nb::cast<absl::string_view>(name_));
-}
-
-std::string GetAttrKey::ToReprString() const {
-  return absl::StrFormat("GetAttrKey(name='%s')",
-                         nb::cast<absl::string_view>(name_));
-}
-
-std::string FlattenedIndexKey::ToString() const {
-  return absl::StrFormat("[<flat index %d>]", key_);
-}
-
-std::string FlattenedIndexKey::ToReprString() const {
-  return absl::StrFormat("FlattenedIndexKey(key=%d)", key_);
-}
-
-bool SequenceKey::Equals(const nb::object& other) {
-  SequenceKey other_key(0);
-  if (!nb::try_cast<SequenceKey>(other, other_key)) return false;
-  return idx_ == other_key.idx();
-}
-
-bool DictKey::Equals(const nb::object& other) {
-  DictKey other_key(nb::none());
-  if (!nb::try_cast<DictKey>(other, other_key)) return false;
-  return key_.equal(other_key.key());
-}
-
-bool GetAttrKey::Equals(const nb::object& other) {
-  GetAttrKey other_key(nb::str(""));
-  if (!nb::try_cast<GetAttrKey>(other, other_key)) return false;
-  return name_.equal(other_key.name());
-}
-
-bool FlattenedIndexKey::Equals(const nb::object& other) {
-  FlattenedIndexKey other_key(0);
-  if (!nb::try_cast<FlattenedIndexKey>(other, other_key)) return false;
-  return key_ == other_key.key();
-}
-
-nanobind::tuple SequenceKey::MatchArgs(nanobind::handle unused) {
-  return nanobind::make_tuple("idx");
-};
-
-nanobind::tuple DictKey::MatchArgs(nanobind::handle unused) {
-  return nanobind::make_tuple("key");
-};
-
-nanobind::tuple GetAttrKey::MatchArgs(nanobind::handle unused) {
-  return nanobind::make_tuple("name");
-};
-
-nanobind::tuple FlattenedIndexKey::MatchArgs(nanobind::handle unused) {
-  return nanobind::make_tuple("key");
-};
-
-template <typename T>
-void PyTreeDef::FlattenImpl(nb::handle handle, T& leaves,
-                            const std::optional<nb::callable>& leaf_predicate,
-                            std::optional<std::vector<nb::object>>& keypath) {
-  Node node;
-  const int start_num_nodes = traversal_.size();
-  const int start_num_leaves = leaves.size();
-  bool is_known_leaf = false;
-  if (leaf_predicate) {
-    nb::object o = (*leaf_predicate)(handle);
-    // Historically we accepted "truthy" values from leaf predicates. Accept
-    // None here to keep existing clients happy.
-    if (o.is_none()) {
-      is_known_leaf = false;
-    } else if (!nb::try_cast<bool>(o, is_known_leaf)) {
-      throw std::invalid_argument(absl::StrCat(
-          "is_leaf predicate returned a non-boolean value ",
-          nb::cast<absl::string_view>(nb::repr(o)), "; expected a boolean"));
-    }
-  }
-  if (is_known_leaf) {
-    nb::object value = nb::borrow<nb::object>(handle);
-    if (keypath.has_value()) {
-      const std::vector<nb::object>& frozen_keypath = keypath.value();
-      nb::object kp_tuple = nb::steal(PyTuple_New(frozen_keypath.size()));
-      for (int i = 0; i < frozen_keypath.size(); ++i) {
-        PyTuple_SET_ITEM(kp_tuple.ptr(), i,
-                         nb::object(frozen_keypath[i]).release().ptr());
-      }
-      value = nb::make_tuple(std::move(kp_tuple), std::move(value));
-    }
-    if constexpr (std::is_same_v<T, nb::list>) {
-      leaves.append(std::move(value));
-    } else {
-      leaves.push_back(std::move(value));
-    }
-  } else {
-    node.kind = registry_->KindOfObject(handle, &node.custom);
-    auto recurse = [this, &leaf_predicate, &leaves](
-                       nb::handle child,
-                       std::optional<std::vector<nb::object>>& keypath) {
-      if (Py_EnterRecursiveCall(
-              " in flatten; PyTree may have cyclical node references.")) {
-        return;
-      }
-      FlattenImpl(child, leaves, leaf_predicate, keypath);
-      Py_LeaveRecursiveCall();
-    };
-    switch (node.kind) {
-      case PyTreeKind::kNone:
-        // Nothing to do.
-        break;
-      case PyTreeKind::kTuple: {
-        node.arity = PyTuple_GET_SIZE(handle.ptr());
-        for (int i = 0; i < node.arity; ++i) {
-          if (keypath.has_value()) {
-            keypath->push_back(make_nb_class<SequenceKey>(i));
-          }
-          recurse(PyTuple_GET_ITEM(handle.ptr(), i), keypath);
-          if (keypath.has_value()) {
-            keypath->pop_back();
-          }
-        }
-        break;
-      }
-      case PyTreeKind::kList: {
-        node.arity = PyList_GET_SIZE(handle.ptr());
-        for (int i = 0; i < node.arity; ++i) {
-          if (keypath.has_value()) {
-            keypath->push_back(make_nb_class<SequenceKey>(i));
-          }
-          recurse(PyList_GET_ITEM(handle.ptr(), i), keypath);
-          if (keypath.has_value()) {
-            keypath->pop_back();
-          }
-        }
-        break;
-      }
-      case PyTreeKind::kDict: {
-        nb::dict dict = nb::borrow<nb::dict>(handle);
-
-        std::vector<nb::object> keys = GetSortedPyDictKeys(dict.ptr());
-        for (nb::object& key : keys) {
-          if (keypath.has_value()) {
-            keypath->push_back(make_nb_class<DictKey>(key));
-          }
-          recurse(dict[key], keypath);
-          if (keypath.has_value()) {
-            keypath->pop_back();
-          }
-        }
-        node.arity = dict.size();
-        node.sorted_dict_keys = std::move(keys);
-        break;
-      }
-      case PyTreeKind::kCustom: {
-        if (keypath.has_value()) {
-          auto [leaves, aux_data] = node.custom->ToIterableWithKeys(handle);
-          node.node_data = std::move(aux_data);
-          node.arity = 0;
-          for (auto& [key, leaf] : leaves) {
-            keypath->push_back(key);
-            ++node.arity;
-            recurse(leaf, keypath);
-            keypath->pop_back();
-          }
-        } else {
-          auto [leaves, aux_data] = node.custom->ToIterable(handle);
-          node.node_data = std::move(aux_data);
-          node.arity = 0;
-          for (nb::handle entry : leaves) {
-            ++node.arity;
-            recurse(entry, keypath);
-          }
-        }
-        break;
-      }
-      case PyTreeKind::kDataclass: {
-        auto meta_size = node.custom->meta_fields.size();
-        nb::object aux_data = nb::steal(PyTuple_New(meta_size));
-        for (int meta_leaf = 0; meta_leaf < meta_size; ++meta_leaf) {
-          PyTuple_SET_ITEM(
-              aux_data.ptr(), meta_leaf,
-              nb::getattr(handle, node.custom->meta_fields[meta_leaf])
-                  .release()
-                  .ptr());
-        }
-        node.node_data = std::move(aux_data);
-        auto data_size = node.custom->data_fields.size();
-        node.arity = data_size;
-        for (int leaf = 0; leaf < data_size; ++leaf) {
-          if (keypath.has_value()) {
-            keypath->push_back(
-                make_nb_class<GetAttrKey>(node.custom->data_fields[leaf]));
-          }
-          recurse(nb::getattr(handle, node.custom->data_fields[leaf]), keypath);
-          if (keypath.has_value()) {
-            keypath->pop_back();
-          }
-        }
-        break;
-      }
-      case PyTreeKind::kNamedTuple: {
-        nb::tuple tuple = nb::borrow<nb::tuple>(handle);
-        node.arity = tuple.size();
-        node.node_data = nb::borrow<nb::object>(tuple.type());
-        if (keypath.has_value()) {
-          // Get key names from NamedTuple fields.
-          nb::tuple fields;
-          if (!nb::try_cast<nb::tuple>(nb::getattr(tuple, "_fields"), fields) ||
-              tuple.size() != fields.size()) {
-            throw std::invalid_argument(
-                "A namedtuple's _fields attribute should have the same size as "
-                "the tuple.");
-          }
-          auto field_iter = fields.begin();
-          for (nb::handle entry : tuple) {
-            keypath->push_back(make_nb_class<GetAttrKey>(nb::str(*field_iter)));
-            field_iter++;
-            recurse(entry, keypath);
-            keypath->pop_back();
-          }
-        } else {
-          for (nb::handle entry : tuple) {
-            recurse(entry, keypath);
-          }
-        }
-        break;
-      }
-      default:
-        DCHECK(node.kind == PyTreeKind::kLeaf);
-        auto value = nb::borrow<nb::object>(handle);
-        if (keypath.has_value()) {
-          const std::vector<nb::object>& frozen_keypath = keypath.value();
-          nb::object kp_tuple = nb::steal(PyTuple_New(frozen_keypath.size()));
-          for (int i = 0; i < frozen_keypath.size(); ++i) {
-            PyTuple_SET_ITEM(kp_tuple.ptr(), i,
-                             nb::object(frozen_keypath[i]).release().ptr());
-          }
-          value = nb::make_tuple(std::move(kp_tuple), std::move(value));
-        }
-        if constexpr (std::is_same_v<T, nb::list>) {
-          leaves.append(std::move(value));
-        } else {
-          leaves.push_back(std::move(value));
-        }
-    }
-  }
-  node.num_nodes = traversal_.size() - start_num_nodes + 1;
-  node.num_leaves = leaves.size() - start_num_leaves;
-  traversal_.push_back(std::move(node));
-}
-
-void PyTreeDef::Flatten(nb::handle handle,
-                        absl::InlinedVector<nb::object, 2>& leaves,
-                        std::optional<nb::callable> leaf_predicate) {
-  std::optional<std::vector<nb::object>> keypath = std::nullopt;
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
-}
-
-void PyTreeDef::Flatten(nb::handle handle, std::vector<nb::object>& leaves,
-                        std::optional<nb::callable> leaf_predicate) {
-  std::optional<std::vector<nb::object>> keypath = std::nullopt;
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
-}
-
-void PyTreeDef::Flatten(nb::handle handle, nb::list& leaves,
-                        std::optional<nb::callable> leaf_predicate) {
-  std::optional<std::vector<nb::object>> keypath = std::nullopt;
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
-}
-
-/*static*/ std::pair<std::vector<nb::object>, nb_class_ptr<PyTreeDef>>
-PyTreeDef::Flatten(nb::handle x, nb_class_ptr<PyTreeRegistry> registry,
-                   std::optional<nb::callable> leaf_predicate) {
-  auto def = make_nb_class<PyTreeDef>(registry);
-  std::vector<nb::object> leaves;
-  def->Flatten(x, leaves, leaf_predicate);
-  return std::make_pair(std::move(leaves), std::move(def));
-}
-
-void PyTreeDef::FlattenWithPath(nb::handle handle, nanobind::list& leaves,
-                                std::optional<nb::callable> leaf_predicate) {
-  std::optional<std::vector<nb::object>> keypath = std::vector<nb::object>();
-  FlattenImpl(handle, leaves, leaf_predicate, keypath);
-}
-
-/*static*/ bool PyTreeDef::AllLeaves(PyTreeRegistry* registry,
-                                     const nb::iterable& x) {
-  const PyTreeRegistry::Registration* custom;
-  for (const nb::handle& h : x) {
-    if (registry->KindOfObject(h, &custom) != PyTreeKind::kLeaf) return false;
-  }
-  return true;
-}
-
-template <typename T>
-nb::object PyTreeDef::UnflattenImpl(T leaves) const {
-  absl::InlinedVector<nb::object, 4> agenda;
-  auto it = leaves.begin();
-  int leaf_count = 0;
-  for (const Node& node : traversal_) {
-    if (agenda.size() < node.arity) {
-      throw std::logic_error("Too few elements for TreeDef node.");
-    }
-    switch (node.kind) {
-      case PyTreeKind::kLeaf:
-        if (it == leaves.end()) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Too few leaves for PyTreeDef; expected %d, got %d", num_leaves(),
-              leaf_count));
-        }
-        agenda.push_back(nb::borrow<nb::object>(*it));
-        ++it;
-        ++leaf_count;
-        break;
-
-      case PyTreeKind::kNone:
-      case PyTreeKind::kTuple:
-      case PyTreeKind::kNamedTuple:
-      case PyTreeKind::kList:
-      case PyTreeKind::kDict:
-      case PyTreeKind::kCustom:
-      case PyTreeKind::kDataclass: {
-        const int size = agenda.size();
-        absl::Span<nb::object> span;
-        if (node.arity > 0) {
-          span = absl::Span<nb::object>(&agenda[size - node.arity], node.arity);
-        }
-        nb::object o = MakeNode(node, span);
-        agenda.resize(size - node.arity);
-        agenda.push_back(o);
-        break;
-      }
-    }
-  }
-  if (it != leaves.end()) {
-    throw std::invalid_argument(absl::StrFormat(
-        "Too many leaves for PyTreeDef; expected %d.", num_leaves()));
-  }
-  if (agenda.size() != 1) {
-    throw std::logic_error("PyTreeDef traversal did not yield a singleton.");
-  }
-  return std::move(agenda.back());
-}
-
-nb::object PyTreeDef::Unflatten(nb::iterable leaves) const {
-  return UnflattenImpl(leaves);
-}
-
-nb::object PyTreeDef::Unflatten(absl::Span<const nb::object> leaves) const {
-  return UnflattenImpl(leaves);
-}
-
-/*static*/ nb::object PyTreeDef::MakeNode(const PyTreeDef::Node& node,
-                                          absl::Span<nb::object> children) {
-  if (children.size() != node.arity) {
-    throw std::logic_error("Node arity mismatch.");
-  }
-  switch (node.kind) {
-    case PyTreeKind::kLeaf:
-      throw std::logic_error("MakeNode not implemented for leaves.");
-
-    case PyTreeKind::kNone:
-      return nb::none();
-
-    case PyTreeKind::kTuple:
-    case PyTreeKind::kNamedTuple: {
-      nb::object tuple = nb::steal(PyTuple_New(node.arity));
-      for (int i = 0; i < node.arity; ++i) {
-        PyTuple_SET_ITEM(tuple.ptr(), i, children[i].release().ptr());
-      }
-      if (node.kind == PyTreeKind::kNamedTuple) {
-        return node.node_data(*tuple);
-      } else {
-        return tuple;
-      }
-    }
-
-    case PyTreeKind::kList: {
-      nb::object list = nb::steal(PyList_New(node.arity));
-      for (int i = 0; i < node.arity; ++i) {
-        PyList_SET_ITEM(list.ptr(), i, children[i].release().ptr());
-      }
-      return list;
-    }
-
-    case PyTreeKind::kDict: {
-      nb::dict dict;
-      for (int i = 0; i < node.arity; ++i) {
-        dict[node.sorted_dict_keys[i]] = std::move(children[i]);
-      }
-      return std::move(dict);
-      break;
-    }
-    case PyTreeKind::kCustom: {
-      nb::object tuple = nb::steal(PyTuple_New(node.arity));
-      for (int i = 0; i < node.arity; ++i) {
-        PyTuple_SET_ITEM(tuple.ptr(), i, children[i].release().ptr());
-      }
-      return node.custom->from_iterable(node.node_data, tuple);
-    }
-
-    case PyTreeKind::kDataclass: {
-      nb::kwargs kwargs;
-      auto meta_size = node.custom->meta_fields.size();
-      for (int i = 0; i < meta_size; ++i) {
-        kwargs[node.custom->meta_fields[i]] =
-            nb::borrow(nb::tuple(node.node_data)[i]);
-      }
-      auto data_size = node.custom->data_fields.size();
-      for (int i = 0; i < data_size; ++i) {
-        kwargs[node.custom->data_fields[i]] = std::move(children[i]);
-      }
-      return node.custom->type(**kwargs);
-    }
-  }
-  throw std::logic_error("Unreachable code.");
-}
-
-nb::list PyTreeDef::FlattenUpTo(nb::handle xs) const {
-  nb::list leaves = nb::steal<nb::list>(PyList_New(num_leaves()));
-  std::vector<nb::object> agenda;
-  agenda.push_back(nb::borrow<nb::object>(xs));
-  auto it = traversal_.rbegin();
-  int leaf = num_leaves() - 1;
-  while (!agenda.empty()) {
-    if (it == traversal_.rend()) {
-      throw std::invalid_argument(absl::StrFormat(
-          "Tree structures did not match: %s vs %s",
-          nb::cast<absl::string_view>(nb::repr(xs)), ToString()));
-    }
-    const Node& node = *it;
-    nb::object object = agenda.back();
-    agenda.pop_back();
-    ++it;
-
-    switch (node.kind) {
-      case PyTreeKind::kLeaf:
-        if (leaf < 0) {
-          throw std::logic_error("Leaf count mismatch.");
-        }
-        PyList_SET_ITEM(leaves.ptr(), leaf, object.release().ptr());
-        --leaf;
-        break;
-
-      case PyTreeKind::kNone:
-        if (!object.is_none()) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Expected None, got %s.\n\n"
-              "In previous releases of JAX, flatten-up-to used to "
-              "consider None to be a tree-prefix of non-None values. To obtain "
-              "the previous behavior, you can usually write:\n"
-              "  jax.tree.map(lambda x, y: None if x is None else f(x, y), a, "
-              "b, is_leaf=lambda x: x is None)",
-              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        break;
-
-      case PyTreeKind::kTuple: {
-        if (!PyTuple_CheckExact(object.ptr())) {
-          throw std::invalid_argument(
-              absl::StrFormat("Expected tuple, got %s.",
-                              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        nb::tuple tuple = nb::borrow<nb::tuple>(object);
-        if (tuple.size() != node.arity) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Tuple arity mismatch: %d != %d; tuple: %s.", tuple.size(),
-              node.arity, nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        for (nb::handle entry : tuple) {
-          agenda.push_back(nb::borrow<nb::object>(entry));
-        }
-        break;
-      }
-
-      case PyTreeKind::kList: {
-        if (!PyList_CheckExact(object.ptr())) {
-          throw std::invalid_argument(
-              absl::StrFormat("Expected list, got %s.",
-                              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        nb::list list = nb::borrow<nb::list>(object);
-        if (list.size() != node.arity) {
-          throw std::invalid_argument(absl::StrFormat(
-              "List arity mismatch: %d != %d; list: %s.", list.size(),
-              node.arity, nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        for (nb::handle entry : list) {
-          agenda.push_back(nb::borrow<nb::object>(entry));
-        }
-        break;
-      }
-
-      case PyTreeKind::kDict: {
-        if (!PyDict_CheckExact(object.ptr())) {
-          throw std::invalid_argument(
-              absl::StrFormat("Expected dict, got %s.",
-                              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        nb::dict dict = nb::borrow<nb::dict>(object);
-        std::vector<nb::object> keys = GetSortedPyDictKeys(dict.ptr());
-        if (!IsSortedPyDictKeysEqual(keys, node.sorted_dict_keys)) {
-          // Convert to a nb::list for nb::repr to avoid having to stringify a
-          // vector. This is error path so it is fine to pay conversion cost.
-          throw std::invalid_argument(
-              absl::StrFormat("Dict key mismatch; expected keys: %s; dict: %s.",
-                              nb::cast<absl::string_view>(
-                                  nb::repr(nb::cast(node.sorted_dict_keys))),
-                              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        for (nb::handle key : keys) {
-          agenda.push_back(dict[key]);
-        }
-        break;
-      }
-
-      case PyTreeKind::kNamedTuple: {
-        if (!nb::isinstance<nb::tuple>(object) ||
-            !nb::hasattr(object, "_fields")) {
-          throw std::invalid_argument(
-              absl::StrFormat("Expected named tuple, got %s.",
-                              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        nb::tuple tuple = nb::borrow<nb::tuple>(object);
-        if (tuple.size() != node.arity) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Named tuple arity mismatch: %d != %d; tuple: %s.", tuple.size(),
-              node.arity, nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        if (tuple.type().not_equal(node.node_data)) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Named tuple type mismatch: expected type: %s, tuple: %s.",
-              nb::cast<absl::string_view>(nb::repr(node.node_data)),
-              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        for (nb::handle entry : tuple) {
-          agenda.push_back(nb::borrow<nb::object>(entry));
-        }
-        break;
-      }
-
-      case PyTreeKind::kCustom: {
-        auto* registration = registry_->Lookup(object.type());
-        if (registration != node.custom) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Custom node type mismatch: expected type: %s, value: %s.",
-              nb::cast<absl::string_view>(nb::repr(node.custom->type)),
-              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        auto [leaves, aux_data] = node.custom->ToIterable(object);
-        if (node.node_data.not_equal(aux_data)) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Mismatch custom node data: %s != %s; value: %s.",
-              nb::cast<absl::string_view>(nb::repr(node.node_data)),
-              nb::cast<absl::string_view>(nb::repr(aux_data)),
-              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        int arity = 0;
-        for (nb::handle entry : leaves) {
-          ++arity;
-          agenda.push_back(nb::borrow<nb::object>(entry));
-        }
-        if (arity != node.arity) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Custom type arity mismatch: %d != %d; value: %s.", arity,
-              node.arity, nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        break;
-      }
-
-      case PyTreeKind::kDataclass: {
-        auto* registration = registry_->Lookup(object.type());
-        if (registration != node.custom) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Custom dataclasss node type mismatch: expected type: %s, value: "
-              "%s.",
-              nb::cast<absl::string_view>(nb::repr(node.custom->type)),
-              nb::cast<absl::string_view>(nb::repr(std::move(object)))));
-        }
-        auto meta_size = node.custom->meta_fields.size();
-        nb::object aux_data = nb::steal(PyTuple_New(meta_size));
-        for (int meta_leaf = 0; meta_leaf < meta_size; ++meta_leaf) {
-          PyTuple_SET_ITEM(
-              aux_data.ptr(), meta_leaf,
-              nb::getattr(object, node.custom->meta_fields[meta_leaf])
-                  .release()
-                  .ptr());
-        }
-        if (node.node_data.not_equal(aux_data)) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Mismatch custom dataclass node data: %s != %s; value: %s.",
-              nb::cast<absl::string_view>(nb::repr(node.node_data)),
-              nb::cast<absl::string_view>(nb::repr(aux_data)),
-              nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        auto data_size = node.custom->data_fields.size();
-        if (data_size != node.arity) {
-          throw std::invalid_argument(absl::StrFormat(
-              "Custom type arity mismatch: %d != %d; value: %s.", data_size,
-              node.arity, nb::cast<absl::string_view>(nb::repr(object))));
-        }
-        for (int leaf = 0; leaf < data_size; ++leaf) {
-          agenda.push_back(nb::borrow<nb::object>(
-              nb::getattr(object, node.custom->data_fields[leaf])));
-        }
-        break;
-      }
-    }
-  }
-  if (it != traversal_.rend() || leaf != -1) {
-    throw std::invalid_argument(
-        absl::StrFormat("Tree structures did not match: %s vs %s",
-                        nb::cast<absl::string_view>(nb::repr(xs)), ToString()));
-  }
-  return leaves;
-}
-
-nb::object PyTreeDef::Walk(const nb::callable& f_node, nb::handle f_leaf,
-                           nb::iterable leaves) const {
-  std::vector<nb::object> agenda;
-  auto it = leaves.begin();
-  for (const Node& node : traversal_) {
-    switch (node.kind) {
-      case PyTreeKind::kLeaf: {
-        if (it == leaves.end()) {
-          throw std::invalid_argument("Too few leaves for PyTreeDef");
-        }
-
-        nb::object leaf = nb::borrow<nb::object>(*it);
-        agenda.push_back(f_leaf.is_none() ? std::move(leaf)
-                                          : f_leaf(std::move(leaf)));
-        ++it;
-        break;
-      }
-
-      case PyTreeKind::kNone:
-      case PyTreeKind::kTuple:
-      case PyTreeKind::kNamedTuple:
-      case PyTreeKind::kList:
-      case PyTreeKind::kDict:
-      case PyTreeKind::kCustom:
-      case PyTreeKind::kDataclass: {
-        if (agenda.size() < node.arity) {
-          throw std::logic_error("Too few elements for custom type.");
-        }
-        nb::object tuple = nb::steal(PyTuple_New(node.arity));
-        for (int i = node.arity - 1; i >= 0; --i) {
-          PyTuple_SET_ITEM(tuple.ptr(), i, agenda.back().release().ptr());
-          agenda.pop_back();
-        }
-        nb::object node_data = node.node_data;
-        if (node.kind == PyTreeKind::kDict) {
-          // Convert to a nb::list for f_node invocation.
-          node_data = nb::cast(node.sorted_dict_keys);
-        }
-        agenda.push_back(f_node(tuple, node_data ? node_data : nb::none()));
-      }
-    }
-  }
-  if (it != leaves.end()) {
-    throw std::invalid_argument("Too many leaves for PyTreeDef");
-  }
-  if (agenda.size() != 1) {
-    throw std::logic_error("PyTreeDef traversal did not yield a singleton.");
-  }
-  return std::move(agenda.back());
-}
-
-nb::object PyTreeDef::FromIterableTreeHelper(
-    nb::handle xs,
-    absl::InlinedVector<PyTreeDef::Node, 1>::const_reverse_iterator* it) const {
-  if (*it == traversal_.rend()) {
-    throw std::invalid_argument("Tree structures did not match.");
-  }
-  const Node& node = **it;
-  ++*it;
-  if (node.kind == PyTreeKind::kLeaf) {
-    return nb::borrow<nb::object>(xs);
-  }
-  nb::iterable iterable = nb::borrow<nb::iterable>(xs);
-  std::vector<nb::object> ys;
-  ys.reserve(node.arity);
-  for (nb::handle x : iterable) {
-    ys.push_back(nb::borrow<nb::object>(x));
-  }
-  if (ys.size() != node.arity) {
-    throw std::invalid_argument("Arity mismatch between trees");
-  }
-  for (int j = node.arity - 1; j >= 0; --j) {
-    ys[j] = FromIterableTreeHelper(ys[j], it);
-  }
-
-  return MakeNode(node, absl::MakeSpan(ys));
-}
-
-nb::object PyTreeDef::FromIterableTree(nb::handle xs) const {
-  auto it = traversal_.rbegin();
-  nb::object out = FromIterableTreeHelper(xs, &it);
-  if (it != traversal_.rend()) {
-    throw std::invalid_argument("Tree structures did not match.");
-  }
-  return out;
-}
-
-nb_class_ptr<PyTreeDef> PyTreeDef::Compose(const PyTreeDef& inner) const {
-  if (inner.registry_ != registry_) {
-    throw std::invalid_argument(
-        "PyTree registries of PyTreeDefs passed to Compose() must match.");
-  }
-  auto out = make_nb_class<PyTreeDef>(registry_ref_);
-  out->traversal_.reserve(static_cast<size_t>(num_leaves()) *
-                              inner.num_nodes() +
-                          num_nodes() - num_leaves());
-  for (const Node& n : traversal_) {
-    if (n.kind == PyTreeKind::kLeaf) {
-      absl::c_copy(inner.traversal_, std::back_inserter(out->traversal_));
-    } else {
-      out->traversal_.push_back(n);
-    }
-  }
-  out->SetNumLeavesAndNumNodes();
-  return out;
-}
-
-/*static*/ nb_class_ptr<PyTreeDef> PyTreeDef::Tuple(
-    nb_class_ptr<PyTreeRegistry> registry, nb::list defs) {
-  auto out = make_nb_class<PyTreeDef>(std::move(registry));
-  int num_leaves = 0;
-  for (nb::handle def_handle : defs) {
-    const PyTreeDef* def = nb::cast<const PyTreeDef*>(def_handle);
-    if (def->registry() != out->registry()) {
-      throw std::invalid_argument(
-          "PyTree registries of PyTreeDefs passed to Tuple() must match.");
-    }
-    absl::c_copy(def->traversal_, std::back_inserter(out->traversal_));
-    num_leaves += def->num_leaves();
-  }
-  Node node;
-  node.kind = PyTreeKind::kTuple;
-  node.arity = defs.size();
-  node.num_leaves = num_leaves;
-  node.num_nodes = out->traversal_.size() + 1;
-  out->traversal_.push_back(node);
-  return out;
-}
-
-std::vector<nb_class_ptr<PyTreeDef>> PyTreeDef::Children() const {
-  std::vector<nb_class_ptr<PyTreeDef>> children;
-  if (traversal_.empty()) {
-    return children;
-  }
-  Node const& root = traversal_.back();
-  children.resize(root.arity);
-  int pos = traversal_.size() - 1;
-  for (int i = root.arity - 1; i >= 0; --i) {
-    children[i] = make_nb_class<PyTreeDef>(registry_ref_);
-    const Node& node = traversal_.at(pos - 1);
-    if (pos < node.num_nodes) {
-      throw std::logic_error("children() walked off start of array");
-    }
-    std::copy(traversal_.begin() + pos - node.num_nodes,
-              traversal_.begin() + pos,
-              std::back_inserter(children[i]->traversal_));
-    pos -= node.num_nodes;
-  }
-  if (pos != 0) {
-    throw std::logic_error("pos != 0 at end of PyTreeDef::Children");
-  }
-  return children;
-}
-
-std::string PyTreeDef::ToString() const {
-  std::vector<std::string> agenda;
-  for (const Node& node : traversal_) {
-    if (agenda.size() < node.arity) {
-      throw std::logic_error("Too few elements for container.");
-    }
-
-    std::string children =
-        absl::StrJoin(agenda.end() - node.arity, agenda.end(), ", ");
-    std::string representation;
-    switch (node.kind) {
-      case PyTreeKind::kLeaf:
-        agenda.push_back("*");
-        continue;
-      case PyTreeKind::kNone:
-        representation = "None";
-        break;
-      case PyTreeKind::kTuple:
-        // Tuples with only one element must have a trailing comma.
-        if (node.arity == 1) children += ",";
-        representation = absl::StrCat("(", children, ")");
-        break;
-      case PyTreeKind::kList:
-        representation = absl::StrCat("[", children, "]");
-        break;
-      case PyTreeKind::kDict: {
-        if (node.sorted_dict_keys.size() != node.arity) {
-          throw std::logic_error("Number of keys and entries does not match.");
-        }
-        representation = "{";
-        std::string separator;
-        auto child_iter = agenda.end() - node.arity;
-        for (const nb::handle& key : node.sorted_dict_keys) {
-          absl::StrAppendFormat(&representation, "%s%s: %s", separator,
-                                nb::cast<absl::string_view>(nb::repr(key)),
-                                *child_iter);
-          child_iter++;
-          separator = ", ";
-        }
-        representation += "}";
-        break;
-      }
-
-      case PyTreeKind::kNamedTuple:
-      case PyTreeKind::kCustom:
-      case PyTreeKind::kDataclass: {
-        std::string kind;
-        std::string data;
-        if (node.kind == PyTreeKind::kNamedTuple) {
-          kind = "namedtuple";
-          if (node.node_data) {
-            // Node data for named tuples is the type.
-            data = absl::StrFormat(
-                "[%s]", nb::cast<absl::string_view>(
-                            nb::str(nb::getattr(node.node_data, "__name__"))));
-          }
-        } else {
-          kind = nb::cast<std::string>(
-              nb::str(nb::getattr(node.custom->type, "__name__")));
-          if (node.node_data) {
-            data = absl::StrFormat(
-                "[%s]", nb::cast<absl::string_view>(nb::str(node.node_data)));
-          }
-        }
-
-        representation =
-            absl::StrFormat("CustomNode(%s%s, [%s])", kind, data, children);
-        break;
-      }
-    }
-    agenda.erase(agenda.end() - node.arity, agenda.end());
-    agenda.push_back(std::move(representation));
-  }
-  if (agenda.size() != 1) {
-    throw std::logic_error("PyTreeDef traversal did not yield a singleton.");
-  }
-  return absl::StrCat("PyTreeDef(", agenda.back(), ")");
-}
-
-nb::object PyTreeDef::ToPickle() const {
-  nb::list traversal;
-  for (const auto& node : traversal_) {
-    nb::object node_data = node.node_data;
-    if (node.kind == PyTreeKind::kDict) {
-      // Convert to a nb::list for pickling to avoid having to pickle a vector.
-      // Pickle should be a rare operation so this conversion cost is hopefully
-      // on non-critical path.
-      node_data = nb::cast(node.sorted_dict_keys);
-    }
-    traversal.append(
-        nb::make_tuple(static_cast<int>(node.kind), node.arity,
-                       node_data ? node_data : nb::none(),
-                       node.custom != nullptr ? node.custom->type : nb::none(),
-                       node.num_leaves, node.num_nodes));
-  }
-  return nb::make_tuple(nb::cast(registry_ref_), traversal);
-}
-
-void PyTreeDef::FromPickle(nb::object pickle) {
-  for (const auto& item : nb::cast<nb::list>(pickle)) {
-    auto t = nb::cast<nb::tuple>(item);
-    if (t.size() != 6) {
-      throw xla::XlaRuntimeError("Malformed pickled PyTreeDef");
-    }
-    Node& node = traversal_.emplace_back();
-    node.kind = static_cast<PyTreeKind>(nb::cast<int>(t[0]));
-    node.arity = nb::cast<int>(t[1]);
-    switch (node.kind) {
-      case PyTreeKind::kNamedTuple:
-        node.node_data = t[2];
-        break;
-      case PyTreeKind::kDict:
-        node.sorted_dict_keys = nb::cast<std::vector<nb::object>>(t[2]);
-        break;
-      case PyTreeKind::kCustom:
-      case PyTreeKind::kDataclass:
-        node.node_data = t[2];
-        break;
-      default:
-        if (!t[2].is_none()) {
-          throw xla::XlaRuntimeError("Malformed pickled PyTreeDef");
-        }
-        break;
-    }
-    if (node.kind == PyTreeKind::kCustom ||
-        node.kind == PyTreeKind::kDataclass) {
-      node.custom = t[3].is_none() ? nullptr : registry()->Lookup(t[3]);
-      if (node.custom == nullptr) {
-        throw xla::XlaRuntimeError(
-            absl::StrCat("Unknown custom type in pickled PyTreeDef: ",
-                         nb::cast<absl::string_view>(nb::repr(t[3]))));
-      }
-    } else {
-      if (!t[3].is_none()) {
-        throw xla::XlaRuntimeError("Malformed pickled PyTreeDef");
-      }
-    }
-    node.num_leaves = nb::cast<int>(t[4]);
-    node.num_nodes = nb::cast<int>(t[5]);
-  }
-}
-
-void PyTreeDef::SetNumLeavesAndNumNodes() {
-  // num_leaves and num_nodes are fully determined by arity.
-  std::vector<std::pair<int, int>> starts;
-  int num_leaves = 0;
-  for (int i = 0; i < traversal_.size(); ++i) {
-    std::pair<int, int> start = {num_leaves, i};
-    if (traversal_[i].kind == PyTreeKind::kLeaf) {
-      num_leaves += 1;
-    }
-    if (traversal_[i].arity == 0) {
-      starts.push_back(start);
-    } else {
-      starts.resize(starts.size() - (traversal_[i].arity - 1));
-    }
-    traversal_[i].num_leaves = num_leaves - starts.back().first;
-    traversal_[i].num_nodes = i + 1 - starts.back().second;
-  }
-}
-
-void PyTreeDef::SerializeTo(jax::PyTreeDefProto& result) const {
-  absl::flat_hash_map<std::string, uint32_t> interned_strings;
-  auto intern_str = [&](const std::string& key) {
-    auto [it, added] =
-        interned_strings.emplace(key, result.interned_strings_size());
-    if (added) {
-      result.add_interned_strings(key);
-    }
-    return it->second;
-  };
-  for (const auto& node : traversal_) {
-    auto* node_data = result.add_nodes();
-    node_data->set_arity(node.arity);
-    switch (node.kind) {
-      case PyTreeKind::kLeaf:
-        node_data->set_type(jax::PyTreeNodeType::PY_TREE_KIND_LEAF);
-        break;
-      case PyTreeKind::kList:
-        node_data->set_type(jax::PyTreeNodeType::PY_TREE_KIND_LIST);
-        break;
-      case PyTreeKind::kNone:
-        node_data->set_type(jax::PyTreeNodeType::PY_TREE_KIND_NONE);
-        break;
-      case PyTreeKind::kTuple:
-        node_data->set_type(jax::PyTreeNodeType::PY_TREE_KIND_TUPLE);
-        break;
-      case PyTreeKind::kDict:
-        node_data->set_type(jax::PyTreeNodeType::PY_TREE_KIND_DICT);
-        for (auto& key : node.sorted_dict_keys) {
-          if (!nb::isinstance<nb::str>(key)) {
-            throw std::invalid_argument(
-                "Only string keys are supported in proto pytree "
-                "serialization.");
-          }
-          node_data->mutable_dict_keys()->add_str_id(
-              intern_str(nb::cast<std::string>(key)));
-        }
-        break;
-      default:
-        throw std::invalid_argument(
-            "User-defined nodes are not supported when serializing pytrees as "
-            "protocol buffers. You should either convert the user-defined "
-            "nodes to another type or use pickle instead.");
-        break;
-    }
-  }
-}
-
-nb_class_ptr<PyTreeDef> PyTreeDef::DeserializeFrom(
-    nb_class_ptr<PyTreeRegistry> registry, const jax::PyTreeDefProto& input) {
-  std::vector<nb::object> interned_strings;
-  interned_strings.reserve(input.interned_strings().size());
-  for (auto& s : input.interned_strings()) {
-    interned_strings.push_back(nb::cast(s));
-  }
-  nb_class_ptr<PyTreeDef> result =
-      make_nb_class<PyTreeDef>(std::move(registry));
-  for (auto& node_proto : input.nodes()) {
-    result->traversal_.emplace_back();
-    auto& node = result->traversal_.back();
-    node.arity = node_proto.arity();
-    node.custom = nullptr;
-    switch (node_proto.type()) {
-      case jax::PyTreeNodeType::PY_TREE_KIND_LEAF:
-        node.kind = PyTreeKind::kLeaf;
-        break;
-      case jax::PyTreeNodeType::PY_TREE_KIND_LIST:
-        node.kind = PyTreeKind::kList;
-        break;
-      case jax::PyTreeNodeType::PY_TREE_KIND_NONE:
-        node.kind = PyTreeKind::kNone;
-        break;
-      case jax::PyTreeNodeType::PY_TREE_KIND_TUPLE:
-        node.kind = PyTreeKind::kTuple;
-        break;
-      case jax::PyTreeNodeType::PY_TREE_KIND_DICT:
-        node.kind = PyTreeKind::kDict;
-        for (uint32_t str_id : node_proto.dict_keys().str_id()) {
-          if (str_id >= interned_strings.size()) {
-            throw std::invalid_argument(
-                "Malformed pytree proto (dict_key out of range).");
-          }
-          node.sorted_dict_keys.push_back(interned_strings.at(str_id));
-        }
-        break;
-      default:
-        throw std::invalid_argument(
-            "Malformed pytree proto (invalid node type)");
-        break;
-    }
-  }
-  result->SetNumLeavesAndNumNodes();
-  return result;
-}
-
-std::optional<std::pair<nb::object, nb::object>> PyTreeDef::GetNodeData()
-    const {
-  if (traversal_.empty()) {
-    throw std::logic_error("empty PyTreeDef traversal.");
-  }
-  auto builtin_type = [](PyTypeObject* type_obj) {
-    return nb::borrow<nb::object>(reinterpret_cast<PyObject*>(type_obj));
-  };
-  const auto& node = traversal_.back();
-  switch (node.kind) {
-    case PyTreeKind::kLeaf:
-      return std::nullopt;
-    case PyTreeKind::kNone:
-      return std::make_pair(builtin_type(Py_TYPE(Py_None)), nb::none());
-    case PyTreeKind::kTuple:
-      return std::make_pair(builtin_type(&PyTuple_Type), nb::none());
-    case PyTreeKind::kList:
-      return std::make_pair(builtin_type(&PyList_Type), nb::none());
-    case PyTreeKind::kDict:
-      return std::make_pair(builtin_type(&PyDict_Type),
-                            nb::cast(node.sorted_dict_keys));
-    case PyTreeKind::kNamedTuple:
-      return std::make_pair(node.node_data, nb::none());
-    case PyTreeKind::kCustom:
-    case PyTreeKind::kDataclass:
-      return std::make_pair(node.custom->type, node.node_data);
-  }
-}
-
-nb_class_ptr<PyTreeDef> PyTreeDef::MakeFromNodeDataAndChildren(
-    nb_class_ptr<PyTreeRegistry> registry,
-    std::optional<std::pair<nb::object, nb::object>> node_data,
-    nb::iterable children) {
-  nb_class_ptr<PyTreeDef> result =
-      make_nb_class<PyTreeDef>(std::move(registry));
-  int num_leaves = 0;
-  int arity = 0;
-  for (nb::handle pchild : children) {
-    const PyTreeDef& child = nb::cast<const PyTreeDef&>(pchild);
-    absl::c_copy(child.traversal_, std::back_inserter(result->traversal_));
-    num_leaves += child.num_leaves();
-    ++arity;
-  }
-  result->traversal_.emplace_back();
-  auto& node = result->traversal_.back();
-  node.arity = arity;
-  node.custom = nullptr;
-  node.num_leaves = num_leaves;
-  node.num_nodes = result->traversal_.size();
-  if (node_data == std::nullopt) {
-    node.kind = PyTreeKind::kLeaf;
-    ++node.num_leaves;
-    return result;
-  }
-  int is_nt = PyObject_IsSubclass(node_data->first.ptr(),
-                                  reinterpret_cast<PyObject*>(&PyTuple_Type));
-  if (is_nt == -1) {
-    throw nb::python_error();
-  }
-  if (is_nt != 0 && nb::hasattr(node_data->first, "_fields")) {
-    node.kind = PyTreeKind::kNamedTuple;
-    node.node_data = node_data->first;
-    return result;
-  }
-  auto* registration = result->registry()->Lookup(node_data->first);
-  if (registration == nullptr) {
-    throw std::logic_error(absl::StrFormat(
-        "Could not find type: %s.",
-        nb::cast<absl::string_view>(nb::repr(node_data->first))));
-  }
-  node.kind = registration->kind;
-  if (node.kind == PyTreeKind::kCustom || node.kind == PyTreeKind::kDataclass) {
-    node.custom = registration;
-    node.node_data = node_data->second;
-  } else if (node.kind == PyTreeKind::kNamedTuple) {
-    node.node_data = node_data->first;
-  } else if (node.kind == PyTreeKind::kDict) {
-    node.sorted_dict_keys =
-        nb::cast<std::vector<nb::object>>(node_data->second);
-  }
-  return result;
-}
-
-int PyTreeDef::Node::tp_traverse(visitproc visit, void* arg) const {
-  Py_VISIT(node_data.ptr());
-  for (const auto& key : sorted_dict_keys) {
-    Py_VISIT(key.ptr());
-  }
-  return 0;
-}
-
-/* static */ int PyTreeDef::tp_traverse(PyObject* self, visitproc visit,
-                                        void* arg) {
-  PyTreeDef* treedef = nb::inst_ptr<PyTreeDef>(self);
-  Py_VISIT(Py_TYPE(self));
-  Py_VISIT(treedef->registry_ref_.ptr());
-  for (const auto& node : treedef->traversal_) {
-    node.tp_traverse(visit, arg);
-  }
-  return 0;
-}
-
-/* static */ int PyTreeDef::tp_clear(PyObject* self) {
-  PyTreeDef* treedef = nb::inst_ptr<PyTreeDef>(self);
-  treedef->registry_ref_.reset();
-  treedef->traversal_.clear();
-  return 0;
-}
-
-/* static */ PyType_Slot PyTreeDef::slots_[] = {
-    {Py_tp_traverse, (void*)PyTreeDef::tp_traverse},
-    {Py_tp_clear, (void*)PyTreeDef::tp_clear},
-    {0, nullptr},
-};
-
-void BuildPytreeSubmodule(nb::module_& m) {
-  nb::module_ pytree = m.def_submodule("pytree", "Python tree library");
-  pytree.attr("version") = nb::int_(3);
-
-  nb::class_<PyTreeDef> treedef(pytree, "PyTreeDef",
-                                nb::type_slots(PyTreeDef::slots_));
-
-  nb::class_<PyTreeRegistry> registry(m, "PyTreeRegistry", nb::dynamic_attr(),
-                                      nb::type_slots(PyTreeRegistry::slots_));
-
-  registry.def(nb::init<bool, bool, bool, bool, bool>(),
-               nb::arg("enable_none") = true, nb::arg("enable_tuple") = true,
-               nb::arg("enable_namedtuple") = true,
-               nb::arg("enable_list") = true, nb::arg("enable_dict") = true);
-  registry.def(
-      "flatten",
-      [](nb_class_ptr<PyTreeRegistry> registry, nb::object x,
-         std::optional<nb::callable> leaf_predicate) {
-        nb::list leaves;
-        nb_class_ptr<PyTreeDef> def =
-            make_nb_class<PyTreeDef>(std::move(registry));
-        def->Flatten(x, leaves, leaf_predicate);
-        return nb::make_tuple(std::move(leaves), std::move(def));
-      },
-      nb::arg("tree").none(), nb::arg("leaf_predicate").none() = std::nullopt);
-  registry.def("flatten_one_level", &PyTreeRegistry::FlattenOneLevel,
-               nb::arg("tree").none());
-  registry.def("flatten_one_level_with_keys",
-               &PyTreeRegistry::FlattenOneLevelWithKeys,
-               nb::arg("tree").none());
-  registry.def(
-      "flatten_with_path",
-      [](nb_class_ptr<PyTreeRegistry> registry, nb::object x,
-         std::optional<nb::callable> leaf_predicate) {
-        nb::list leaves;
-        nb_class_ptr<PyTreeDef> def =
-            make_nb_class<PyTreeDef>(std::move(registry));
-        def->FlattenWithPath(x, leaves, leaf_predicate);
-        return nb::make_tuple(std::move(leaves), std::move(def));
-      },
-      nb::arg("tree").none(), nb::arg("leaf_predicate").none() = std::nullopt);
-  registry.def("register_node", &PyTreeRegistry::Register,
-               nb::arg("type").none(), nb::arg("to_iterable").none(),
-               nb::arg("from_iterable").none(),
-               nb::arg("to_iterable_with_keys").none() = std::nullopt);
-  registry.def("register_dataclass_node", &PyTreeRegistry::RegisterDataclass);
-  registry.def("__reduce__",
-               [](nb::object self) { return self.attr("__name__"); });
-
-  pytree.attr("_default_registry") = make_nb_class<PyTreeRegistry>(
-      /*enable_none=*/true, /*enable_tuple=*/true, /*enable_namedtuple=*/true,
-      /*enable_list=*/true, /*enable_dict*/ true);
-  pytree.def("default_registry",
-             [registry = nb::cast<nb_class_ptr<PyTreeRegistry>>(
-                  pytree.attr("_default_registry"))]() { return registry; });
-
-  pytree.attr("PyTreeRegistry") = m.attr("PyTreeRegistry");
-  pytree.def("tuple", &PyTreeDef::Tuple);
-  pytree.def("all_leaves", &PyTreeDef::AllLeaves);
-
-  treedef.def("unflatten",
-              static_cast<nb::object (PyTreeDef::*)(nb::iterable leaves) const>(
-                  &PyTreeDef::Unflatten));
-  treedef.def("flatten_up_to", &PyTreeDef::FlattenUpTo, nb::arg("tree").none());
-  treedef.def("compose", &PyTreeDef::Compose);
-  treedef.def(
-      "walk", &PyTreeDef::Walk,
-      "Walk pytree, calling f_node(node, node_data) at nodes, and f_leaf "
-      "at leaves",
-      nb::arg("f_node"), nb::arg("f_leaf"), nb::arg("leaves"));
-  treedef.def("from_iterable_tree", &PyTreeDef::FromIterableTree);
-  treedef.def("children", &PyTreeDef::Children);
-  treedef.def_prop_ro("num_leaves", &PyTreeDef::num_leaves);
-  treedef.def_prop_ro("num_nodes", &PyTreeDef::num_nodes);
-  treedef.def("__repr__", &PyTreeDef::ToString);
-  treedef.def("__eq__",
-              [](const PyTreeDef& a, const PyTreeDef& b) { return a == b; });
-  treedef.def("__ne__",
-              [](const PyTreeDef& a, const PyTreeDef& b) { return a != b; });
-  treedef.def("__hash__", [](const PyTreeDef& t) { return absl::HashOf(t); });
-  treedef.def("serialize_using_proto", [](const PyTreeDef& a) {
-    jax::PyTreeDefProto result;
-    a.SerializeTo(result);
-    std::string serialized = result.SerializeAsString();
-    return nb::bytes(serialized.data(), serialized.size());
-  });
-  treedef.def_static(
-      "deserialize_using_proto",
-      [](nb_class_ptr<PyTreeRegistry> registry, nb::bytes data) {
-        jax::PyTreeDefProto input;
-        absl::string_view serialized(data.c_str(), data.size());
-        if (serialized.size() > std::numeric_limits<int>::max()) {
-          throw xla::XlaRuntimeError(
-              "Pytree serialization too large to deserialize.");
-        }
-        if (!input.ParseFromArray(serialized.data(), serialized.size())) {
-          throw xla::XlaRuntimeError("Could not deserialize PyTreeDefProto.");
-        }
-        return PyTreeDef::DeserializeFrom(std::move(registry), input);
-      },
-      nb::arg("registry"), nb::arg("data"));
-  treedef.def("node_data", &PyTreeDef::GetNodeData,
-              "Returns None if a leaf-pytree, else (type, node_data)");
-  treedef.def_static(
-      "make_from_node_data_and_children",
-      &PyTreeDef::MakeFromNodeDataAndChildren, nb::arg("registry"),
-      nb::arg("node_data").none(), nb::arg("children"),
-      "Reconstructs a pytree from `node_data()` and `children()`.");
-  treedef.def("__getstate__", &PyTreeDef::ToPickle);
-  treedef.def("__setstate__", [](PyTreeDef& t, nb::object o) {
-    nb::tuple pickle = nb::cast<nb::tuple>(o);
-    if (pickle.size() != 2) {
-      throw xla::XlaRuntimeError(
-          "Malformed pickled PyTreeDef, expected 2-tuple");
-    }
-    auto registry = nb::cast<nb_class_ptr<PyTreeRegistry>>(pickle[0]);
-    new (&t) PyTreeDef(registry);
-    t.FromPickle(pickle[1]);
-  });
-
-  nb::class_<SequenceKey> sequence_key(pytree, "SequenceKey");
-  sequence_key.def(nb::init<int>(), nb::arg("idx"));
-  sequence_key.def("__str__", &SequenceKey::ToString);
-  sequence_key.def("__repr__", &SequenceKey::ToReprString);
-  sequence_key.def("__eq__", &SequenceKey::Equals);
-  sequence_key.def("__hash__", [](const SequenceKey& key) {
-    return key.idx() + kSequenceKeyHashSalt;
-  });
-  sequence_key.def_prop_ro("idx", &SequenceKey::idx);
-  sequence_key.def_prop_ro_static("__match_args__", &SequenceKey::MatchArgs);
-  sequence_key.def("__getstate__",
-                   [](SequenceKey& key) { return nb::make_tuple(key.idx()); });
-  sequence_key.def("__setstate__",
-                   [](SequenceKey& key, const nb::tuple& state) {
-                     if (state.size() != 1) {
-                       throw xla::XlaRuntimeError(
-                           "Malformed pickled SequenceKey, expected 1-tuple");
-                     }
-                     new (&key) SequenceKey(nb::cast<int>(state[0]));
-                   });
-
-  nb::class_<DictKey> dict_key(pytree, "DictKey",
-                               nb::type_slots(DictKey::slots_));
-  dict_key.def(nb::init<nb::object>(), nb::arg("key"));
-  dict_key.def("__str__", &DictKey::ToString);
-  dict_key.def("__repr__", &DictKey::ToReprString);
-  dict_key.def("__eq__", &DictKey::Equals);
-  dict_key.def("__hash__",
-               [](const DictKey& key) { return nanobind::hash(key.key()); });
-  dict_key.def_prop_ro("key", &DictKey::key);
-  dict_key.def_prop_ro_static("__match_args__", &DictKey::MatchArgs);
-  dict_key.def("__getstate__",
-               [](DictKey& key) { return nb::make_tuple(key.key()); });
-  dict_key.def("__setstate__", [](DictKey& key, const nb::tuple& state) {
-    if (state.size() != 1) {
-      throw xla::XlaRuntimeError("Malformed pickled DictKey, expected 1-tuple");
-    }
-    new (&key) DictKey(nb::cast<nb::object>(state[0]));
-  });
-
-  nb::class_<GetAttrKey> get_attr_key(pytree, "GetAttrKey");
-  get_attr_key.def(nb::init<nb::str>(), nb::arg("name"));
-  get_attr_key.def("__str__", &GetAttrKey::ToString);
-  get_attr_key.def("__repr__", &GetAttrKey::ToReprString);
-  get_attr_key.def("__eq__", &GetAttrKey::Equals);
-  get_attr_key.def("__hash__",
-                   [](const GetAttrKey& key) { return nb::hash(key.name()); });
-  get_attr_key.def_prop_ro("name", &GetAttrKey::name);
-  get_attr_key.def_prop_ro_static("__match_args__", &GetAttrKey::MatchArgs);
-  get_attr_key.def("__getstate__",
-                   [](GetAttrKey& key) { return nb::make_tuple(key.name()); });
-  get_attr_key.def("__setstate__", [](GetAttrKey& key, const nb::tuple& state) {
-    if (state.size() != 1) {
-      throw xla::XlaRuntimeError(
-          "Malformed pickled GetAttrKey, expected 1-tuple");
-    }
-    new (&key) GetAttrKey(nb::str(state[0]));
-  });
-
-  nb::class_<FlattenedIndexKey> flattened_index_key(pytree,
-                                                    "FlattenedIndexKey");
-  flattened_index_key.def(nb::init<int>(), nb::arg("key"));
-  flattened_index_key.def("__str__", &FlattenedIndexKey::ToString);
-  flattened_index_key.def("__repr__", &FlattenedIndexKey::ToReprString);
-  flattened_index_key.def("__eq__", &FlattenedIndexKey::Equals);
-  flattened_index_key.def("__hash__", [](const FlattenedIndexKey& key) {
-    return key.key() + kFlattenedIndexKeyHashSalt;
-  });
-  flattened_index_key.def_prop_ro("key", &FlattenedIndexKey::key);
-  flattened_index_key.def_prop_ro_static("__match_args__",
-                                         &FlattenedIndexKey::MatchArgs);
-  flattened_index_key.def("__getstate__", [](FlattenedIndexKey& key) {
-    return nb::make_tuple(key.key());
-  });
-  flattened_index_key.def(
-      "__setstate__", [](FlattenedIndexKey& key, const nb::tuple& state) {
-        if (state.size() != 1) {
-          throw xla::XlaRuntimeError(
-              "Malformed pickled FlattenedIndexKey, expected 1-tuple");
-        }
-        new (&key) FlattenedIndexKey(nb::cast<int>(state[0]));
-      });
-}
-
-}  // namespace xla
diff --git a/jaxlib/xla/pytree.h b/jaxlib/xla/pytree.h
deleted file mode 100644
index 722fe4116..000000000
--- a/jaxlib/xla/pytree.h
+++ /dev/null
@@ -1,408 +0,0 @@
-/* Copyright 2019 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#ifndef JAXLIB_XLA_PYTREE_H_
-#define JAXLIB_XLA_PYTREE_H_
-
-// See https://jax.readthedocs.io/en/latest/pytrees.html for the documentation
-// about pytree.
-
-#include <Python.h>
-
-#include <cstddef>
-#include <memory>
-#include <optional>
-#include <string>
-#include <utility>
-#include <vector>
-
-// placeholder for index annotation headers
-#include "absl/container/flat_hash_map.h"
-#include "absl/container/inlined_vector.h"
-#include "absl/hash/hash.h"
-#include "absl/types/span.h"
-#include "nanobind/nanobind.h"
-#include "jaxlib/xla/pytree.pb.h"
-#include "xla/python/nb_class_ptr.h"
-
-namespace xla {
-
-enum class PyTreeKind {
-  kLeaf,        // An opaque leaf node
-  kNone,        // None.
-  kTuple,       // A tuple
-  kNamedTuple,  // A collections.namedtuple
-  kList,        // A list
-  kDict,        // A dict
-  kCustom,      // A custom type.
-  kDataclass,   // A dataclass.
-};
-
-// Registry of custom node types.
-class PyTreeRegistry {
- public:
-  PyTreeRegistry(bool enable_none, bool enable_tuple, bool enable_namedtuple,
-                 bool enable_list, bool enable_dict);
-
-  PyTreeRegistry(const PyTreeRegistry&) = delete;
-  PyTreeRegistry(PyTreeRegistry&&) = delete;
-  PyTreeRegistry& operator=(const PyTreeRegistry&) = delete;
-  PyTreeRegistry& operator=(PyTreeRegistry&&) = delete;
-
-  struct Registration {
-    PyTreeKind kind;
-
-    // The following values are populated for custom types.
-    // The Python type object, used to identify the type.
-    nanobind::object type;
-    // A function with signature: object -> (iterable, aux_data)
-    nanobind::callable to_iterable;
-    // A function with signature: (aux_data, iterable) -> object
-    nanobind::callable from_iterable;
-    // A function with signature: (aux_data, iterable(keypath, leaf)) -> object
-    std::optional<nanobind::callable> to_iterable_with_keys;
-
-    // Helper that calls to_iterable and validates that it returns a pair
-    // of an iterable and an aux_data object
-    std::pair<nanobind::iterable, nanobind::object> ToIterable(
-        nanobind::handle o) const;
-    // Helper that calls to_iterable_with_keys and validates that it returns a
-    // pair of an iterable of key-leaf pairs and an aux_data object. If
-    // to_iterable_with_keys is not available, return a dummy key for each leaf,
-    // similar to the current jax.tree_util.FlattenedIndexKey.
-    std::pair<std::vector<std::pair<nanobind::object, nanobind::object>>,
-              nanobind::object>
-    ToIterableWithKeys(nanobind::handle o) const;
-
-    // For dataclasses.
-    std::vector<nanobind::str> data_fields;
-    std::vector<nanobind::str> meta_fields;
-
-    int tp_traverse(visitproc visit, void* arg);
-  };
-
-  // Registers a new custom type. Objects of `type` will be treated as container
-  // node types in PyTrees.
-  void Register(
-      nanobind::object type, nanobind::callable to_iterable,
-      nanobind::callable from_iterable,
-      std::optional<nanobind::callable> to_iterable_with_keys = std::nullopt);
-  // Same, but for dataclasses.
-  void RegisterDataclass(nanobind::object type,
-                         std::vector<nanobind::str> data_fields,
-                         std::vector<nanobind::str> meta_fields);
-
-  // Finds the custom type registration for `type`. Returns nullptr if none
-  // exists.
-  const Registration* Lookup(nanobind::handle type) const;
-
-  PyTreeKind KindOfObject(nanobind::handle obj,
-                          PyTreeRegistry::Registration const** custom) const;
-
-  // Flattens a pytree one level, returning either a tuple of the leaves and
-  // the node data, or None, if the entry is a leaf.
-  nanobind::object FlattenOneLevel(nanobind::handle x) const;
-  // Similar to above but returns a key-leaf pair for each leaf.
-  nanobind::object FlattenOneLevelWithKeys(nanobind::handle x) const;
-  // Underlying implementation of FlattenOneLevel and FlattenOneLevelWithKeys.
-  nanobind::object FlattenOneLevelImpl(nanobind::handle x,
-                                       bool with_keys) const;
-
-  static PyType_Slot slots_[];
-
- private:
-  struct TypeHash {
-    using is_transparent = void;
-    size_t operator()(const nanobind::object& t) const {
-      return absl::HashOf(t.ptr());
-    }
-    size_t operator()(const nanobind::handle& t) const {
-      return absl::HashOf(t.ptr());
-    }
-  };
-  struct TypeEq {
-    using is_transparent = void;
-    bool operator()(const nanobind::object& a,
-                    const nanobind::object& b) const {
-      return a.ptr() == b.ptr();
-    }
-    bool operator()(const nanobind::object& a,
-                    const nanobind::handle& b) const {
-      return a.ptr() == b.ptr();
-    }
-  };
-  mutable nanobind::ft_mutex mu_;
-  absl::flat_hash_map<nanobind::object, std::unique_ptr<Registration>, TypeHash,
-                      TypeEq>
-      registrations_;  // Guarded by mu_
-  bool enable_namedtuple_;
-
-  static int tp_traverse(PyObject* self, visitproc visit, void* arg);
-  static int tp_clear(PyObject* self);
-};
-
-class SequenceKey {
- public:
-  explicit SequenceKey(int idx) : idx_(idx) {};
-  std::string ToReprString() const;
-  std::string ToString() const;
-  bool Equals(const nanobind::object& other);
-  int idx() const { return idx_; }
-  static nanobind::tuple MatchArgs(nanobind::handle unused);
-
- private:
-  int idx_;
-};
-
-class DictKey {
- public:
-  explicit DictKey(nanobind::object key) : key_(key) {};
-  std::string ToReprString() const;
-  std::string ToString() const;
-  bool Equals(const nanobind::object& other);
-  nanobind::object key() const { return key_; }
-  static nanobind::tuple MatchArgs(nanobind::handle unused);
-  static PyType_Slot slots_[];
-
- private:
-  nanobind::object key_;
-  static int tp_traverse(PyObject* self, visitproc visit, void* arg);
-  static int tp_clear(PyObject* self);
-};
-
-class GetAttrKey {
- public:
-  explicit GetAttrKey(nanobind::str name) : name_(name) {};
-  std::string ToReprString() const;
-  std::string ToString() const;
-  bool Equals(const nanobind::object& other);
-  nanobind::str name() const { return name_; }
-  static nanobind::tuple MatchArgs(nanobind::handle unused);
-
- private:
-  nanobind::str name_;
-};
-
-class FlattenedIndexKey {
- public:
-  explicit FlattenedIndexKey(int key) : key_(key) {};
-  std::string ToReprString() const;
-  std::string ToString() const;
-  bool Equals(const nanobind::object& other);
-  int key() const { return key_; }
-  static nanobind::tuple MatchArgs(nanobind::handle unused);
-
- private:
-  int key_;
-};
-
-// A PyTreeDef describes the tree structure of a PyTree. A PyTree is a tree of
-// Python values, where the interior nodes are tuples, lists, dictionaries, or
-// user-defined containers, and the leaves are other objects.
-class PyTreeDef {
- public:
-  // Unowned registry: the registry must remain live at least as long as the
-  // PyTreeDef. It is the caller's responsibility to enforce this.
-  explicit PyTreeDef(PyTreeRegistry* registry) : registry_(registry) {}
-
-  explicit PyTreeDef(nb_class_ptr<PyTreeRegistry> registry)
-      : registry_(registry.get()), registry_ref_(std::move(registry)) {}
-
-  // Flattens a Pytree into a list of leaves and a PyTreeDef.
-  // Returns references to the flattened objects, which might be temporary
-  // objects in the case of custom pytype handlers.
-  static std::pair<std::vector<nanobind::object>, nb_class_ptr<PyTreeDef>>
-  Flatten(nanobind::handle x, nb_class_ptr<PyTreeRegistry> registry,
-          std::optional<nanobind::callable> leaf_predicate = std::nullopt);
-
-  // Flattens a Pytree into a list of `leaves` and a PyTreeDef (this).
-  // `leaves` owns references to the flattened objects, which might be
-  // temporary objects in the case of custom pytype handlers.
-  void Flatten(nanobind::handle handle, std::vector<nanobind::object>& leaves,
-               std::optional<nanobind::callable> leaf_predicate = std::nullopt);
-  void Flatten(nanobind::handle handle,
-               absl::InlinedVector<nanobind::object, 2>& leaves,
-               std::optional<nanobind::callable> leaf_predicate = std::nullopt);
-  void Flatten(nanobind::handle handle, nanobind::list& leaves,
-               std::optional<nanobind::callable> leaf_predicate = std::nullopt);
-
-  void FlattenWithPath(
-      nanobind::handle handle, nanobind::list& leaves,
-      std::optional<nanobind::callable> leaf_predicate = std::nullopt);
-
-  // Tests whether the given list is a flat list of leaves.
-  static bool AllLeaves(PyTreeRegistry* registry, const nanobind::iterable& x);
-
-  // Flattens a Pytree up to this PyTreeDef. 'this' must be a tree prefix of
-  // the tree-structure of 'x'. For example, if we flatten a value
-  // [(1, (2, 3)), {"foo": 4}] with a treedef [(*, *), *], the result is the
-  // list of leaves [1, (2, 3), {"foo": 4}].
-  nanobind::list FlattenUpTo(nanobind::handle x) const;
-
-  // Returns an unflattened PyTree given an iterable of leaves and a PyTreeDef.
-  nanobind::object Unflatten(nanobind::iterable leaves) const;
-  nanobind::object Unflatten(absl::Span<const nanobind::object> leaves) const;
-
-  // Composes two PyTreeDefs, replacing the leaves of this tree with copies of
-  // `inner`. The returned PyTreeDef holds a reference to its registry.
-  nb_class_ptr<PyTreeDef> Compose(const PyTreeDef& inner) const;
-
-  // Makes a Tuple PyTreeDef out of a vector of PyTreeDefs.
-  static nb_class_ptr<PyTreeDef> Tuple(nb_class_ptr<PyTreeRegistry> registry,
-                                       nanobind::list defs);
-
-  // The returned PyTreeDefs hold a reference to the registry.
-  std::vector<nb_class_ptr<PyTreeDef>> Children() const;
-
-  // Maps a function over a PyTree structure, applying f_leaf to each leaf, and
-  // f_node(node, node_data) to each container node.
-  nanobind::object Walk(const nanobind::callable& f_node,
-                        nanobind::handle f_leaf,
-                        nanobind::iterable leaves) const;
-
-  // Given a tree of iterables with the same node/leaf structure as this PyTree,
-  // build the corresponding PyTree.
-  // TODO(phawkins): use flattening everywhere instead and delete this method.
-  nanobind::object FromIterableTree(nanobind::handle xs) const;
-
-  int num_leaves() const {
-    if (traversal_.empty()) {
-      return 0;
-    }
-    return traversal_.back().num_leaves;
-  }
-
-  int num_nodes() const { return traversal_.size(); }
-
-  PyTreeRegistry* registry() const { return registry_; }
-
-  size_t Hash() const;
-
-  bool operator==(const PyTreeDef& other) const;
-  bool operator!=(const PyTreeDef& other) const { return !(*this == other); }
-
-  std::string ToString() const;
-
-  // Transforms the PyTreeDef into a pickleable object. Used to implement
-  // `PyTreeDef.__getstate__`.
-  nanobind::object ToPickle() const;
-
-  // Transforms the object returned by `ToPickleable()` back to PyTreeDef. Used
-  // to implement `PyTreeDef.__setstate__`.
-  void FromPickle(nanobind::object pickleable);
-
-  void SerializeTo(jax::PyTreeDefProto& result) const;
-
-  static nb_class_ptr<PyTreeDef> DeserializeFrom(
-      nb_class_ptr<PyTreeRegistry> registry, const jax::PyTreeDefProto& input);
-
-  std::optional<std::pair<nanobind::object, nanobind::object>> GetNodeData()
-      const;
-
-  static nb_class_ptr<PyTreeDef> MakeFromNodeDataAndChildren(
-      nb_class_ptr<PyTreeRegistry> registry,
-      std::optional<std::pair<nanobind::object, nanobind::object>> node_data,
-      nanobind::iterable children);
-
-  static PyType_Slot slots_[];
-
- private:
-  void SetNumLeavesAndNumNodes();
-
-  struct Node {
-    PyTreeKind kind = PyTreeKind::kLeaf;
-
-    // Arity for non-kLeaf types.
-    int arity = 0;
-
-    // Kind-specific auxiliary data. For a kNamedTuple, contains the tuple type
-    // object. For a kDict, use `sorted_dict_keys` field below. For a kCustom
-    // type, contains the auxiliary data returned by the `to_iterable` function.
-    nanobind::object node_data;
-
-    // Kind-specific auxiliary data specialized for kDict. Use a c++ vector
-    // to hold the sorted dict keys instead of a py::list to avoid creating
-    // a new python list object when flattening kDict. For deeply nested dict,
-    // using c++ vector instead of py::list avoids creating too many python
-    // objects that make python gc sweep slow.
-    std::vector<nanobind::object> sorted_dict_keys;
-
-    // Custom type registration. Must be null for non-custom types.
-    const PyTreeRegistry::Registration* custom = nullptr;
-
-    // Number of leaf nodes in the subtree rooted at this node.
-    int num_leaves = 0;
-
-    // Number of leaf and interior nodes in the subtree rooted at this node.
-    int num_nodes = 0;
-
-    int tp_traverse(visitproc visit, void* arg) const;
-  };
-  template <typename H>
-  friend H AbslHashValue(H h, const Node& n);
-
-  template <typename H>
-  friend H AbslHashValue(H h, const PyTreeDef& t);
-
-  // Helper that manufactures an instance of a node given its children.
-  static nanobind::object MakeNode(const Node& node,
-                                   absl::Span<nanobind::object> children);
-
-  // Recursive helper used to implement FromIterableTree()
-  nanobind::object FromIterableTreeHelper(
-      nanobind::handle xs,
-      absl::InlinedVector<PyTreeDef::Node, 1>::const_reverse_iterator* it)
-      const;
-
-  template <typename T>
-  void FlattenImpl(nanobind::handle handle, T& leaves,
-                   const std::optional<nanobind::callable>& leaf_predicate,
-                   std::optional<std::vector<nanobind::object>>& keypath);
-
-  template <typename T>
-  nanobind::object UnflattenImpl(T leaves) const;
-
-  static int tp_traverse(PyObject* self, visitproc visit, void* arg);
-  static int tp_clear(PyObject* self);
-
-  // Pytree registry. Not owned.
-  PyTreeRegistry* registry_;
-  // If this class holds a reference to `registry`, it is held by
-  // `registry_ref_`.
-  nb_class_ptr<PyTreeRegistry> registry_ref_;
-
-  // Nodes, in a post-order traversal. We use an ordered traversal to minimize
-  // allocations, and post-order corresponds to the order we need to rebuild the
-  // tree structure.
-  absl::InlinedVector<Node, 1> traversal_;
-};
-
-template <typename H>
-H AbslHashValue(H h, const PyTreeDef::Node& n) {
-  h = H::combine(std::move(h), n.kind, n.arity, n.custom);
-  return h;
-}
-
-template <typename H>
-H AbslHashValue(H h, const PyTreeDef& t) {
-  h = H::combine(std::move(h), t.traversal_);
-  return h;
-}
-
-void BuildPytreeSubmodule(nanobind::module_& m);
-
-}  // namespace xla
-
-#endif  // JAXLIB_XLA_PYTREE_H_
diff --git a/jaxlib/xla/pytree.proto b/jaxlib/xla/pytree.proto
deleted file mode 100644
index 73c087ef5..000000000
--- a/jaxlib/xla/pytree.proto
+++ /dev/null
@@ -1,32 +0,0 @@
-syntax = "proto3";
-
-package jax;
-
-enum PyTreeNodeType {
-  PY_TREE_KIND_INVALID = 0;
-  PY_TREE_KIND_LEAF = 1;
-  PY_TREE_KIND_LIST = 2;
-  PY_TREE_KIND_NONE = 3;
-  PY_TREE_KIND_TUPLE = 4;
-  PY_TREE_KIND_DICT = 5;
-}
-
-message DictKeysProto {
-  repeated uint32 str_id = 1;
-}
-
-message PyTreeNodeDefProto {
-  // Recovers the tree structure.
-  uint32 arity = 1;
-  // Node type.
-  PyTreeNodeType type = 2;
-  // Only set when type == DICT.
-  DictKeysProto dict_keys = 3;
-}
-
-// A Pytree.
-message PyTreeDefProto {
-  repeated PyTreeNodeDefProto nodes = 1;
-  // Extra strings.
-  repeated string interned_strings = 2;
-}
diff --git a/jaxlib/xla/xla.cc b/jaxlib/xla/xla.cc
index bd3ed3205..fdd4456b2 100644
--- a/jaxlib/xla/xla.cc
+++ b/jaxlib/xla/xla.cc
@@ -91,7 +91,6 @@ limitations under the License.
 #include "jaxlib/xla/mlir.h"
 #include "jaxlib/xla/pjit.h"
 #include "jaxlib/xla/pmap_lib.h"
-#include "jaxlib/xla/pytree.h"
 #include "jaxlib/xla/weakref_lru_cache.h"
 #include "jaxlib/xla/xla_compiler.h"
 #include "xla/pjrt/distributed/key_value_store_interface.h"
@@ -120,6 +119,7 @@ limitations under the License.
 #include "xla/python/py_executable.h"
 #include "xla/python/py_memory_space.h"
 #include "xla/python/python_ref_manager.h"
+#include "xla/python/pytree.h"
 #include "xla/python/sharding.h"
 #include "xla/python/traceback.h"
 #include "xla/tsl/distributed_runtime/preemption/preemption_sync_manager.h"


Test patch:
diff --git a/tests/attrs_test.py b/tests/attrs_test.py
index 2334a7b98..169df3712 100644
--- a/tests/attrs_test.py
+++ b/tests/attrs_test.py
@@ -360,6 +360,49 @@ class AttrsTest(jtu.JaxTestCase):
       return i + 1, None
     _, _ = jax.lax.scan(body, 0, None, length=3)  # don't crash
 
+  @parameterized.parameters([True, False])
+  def test_setattr_doesnt_exist(self, jit):
+    class Thing:
+      ...
+    thing = Thing()
+
+    def f(x):
+      assert (not jit) or tracing_is_ok
+      jax_setattr(thing, 'x', x)
+
+    if jit:
+      f = jax.jit(f)
+
+    tracing_is_ok = True
+    self.assertFalse(hasattr(thing, 'x'))
+    f(1.0)
+    self.assertEqual(thing.x, 1.0)
+    f(2.0)
+    self.assertEqual(thing.x, 2.0)
+
+    tracing_is_ok = False
+    f(3.0)
+    self.assertEqual(thing.x, 3.0)
+
+    del thing.x
+    f(4.0)
+    self.assertEqual(thing.x, 4.0)
+
+    tracing_is_ok = True
+    f(5)
+    self.assertEqual(thing.x, 5)
+
+  def test_setattr_doesnt_exist_doesnt_leave_sentinel_around(self):
+    class Thing:
+      ...
+    thing = Thing()
+
+    def f(x):
+      jax_setattr(thing, 'x', x)
+
+    jax.make_jaxpr(f)(3.)
+    self.assertFalse(hasattr(thing, 'x'))
+
 
 class AttrsJVPTest(jtu.JaxTestCase):