[jaxlib] Revived PyTreeDef::FromNodeDataAndChildren

superbobry · Google-ML-Automation · commit b4f26763ce24 · 2025-10-07T13:39:51.000-07:00
This method is useful for serialization as discussed in #32186. Reverts 0e7c96a PiperOrigin-RevId: 816358350
diff --git a/jaxlib/_jax/pytree.pyi b/jaxlib/_jax/pytree.pyi
@@ -121,6 +121,12 @@ class PyTreeDef:
   def from_iterable_tree(self, __xs: Any): ...
   def node_data(self) -> Tuple[type, Any] | None: ...
   def children(self) -> list[PyTreeDef]: ...
+  @staticmethod
+  def from_node_data_and_children(
+      registry: PyTreeRegistry,
+      node_data: Tuple[type, Any] | None,
+      children: Iterable[PyTreeDef],
+  ) -> PyTreeDef: ...
 
   num_leaves: int
   num_nodes: int
diff --git a/jaxlib/pytree.cc b/jaxlib/pytree.cc
@@ -1550,6 +1550,60 @@ std::optional<std::pair<nb::object, nb::object>> PyTreeDef::GetNodeData()
   }
 }
 
+nb_class_ptr<PyTreeDef> PyTreeDef::FromNodeDataAndChildren(
+    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) {
@@ -1789,6 +1843,21 @@ void BuildPytreeSubmodule(nb::module_& m) {
   treedef.def("node_data", &PyTreeDef::GetNodeData,
               "Returns None if a leaf-pytree, else (type, node_data)",
               nb::sig("def node_data(self) -> tuple[type, Any] | None"));
+  treedef.def_static(
+      "from_node_data_and_children",
+      &PyTreeDef::FromNodeDataAndChildren, nb::arg("registry"),
+      nb::arg("node_data").none(), nb::arg("children"),
+      "Reconstructs a pytree from `node_data()` and `children()`.",
+      nb::sig(
+          // clang-format off
+        "def from_node_data_and_children("
+        "self, "
+        "registry: PyTreeRegistry, "
+        "node_data: tuple[type, Any] | None, "
+        "children: typing.Iterable[PyTreeDef]"
+        ") -> PyTreeDef"
+        // clang-format on
+      ));
   treedef.def("__getstate__", &PyTreeDef::ToPickle);
   treedef.def("__setstate__", [](PyTreeDef& t, nb::object o) {
     nb::tuple pickle = nb::cast<nb::tuple>(o);
diff --git a/jaxlib/pytree.h b/jaxlib/pytree.h
@@ -311,6 +311,11 @@ class PyTreeDef {
   std::optional<std::pair<nanobind::object, nanobind::object>> GetNodeData()
       const;
 
+  static nb_class_ptr<PyTreeDef> FromNodeDataAndChildren(
+      nb_class_ptr<PyTreeRegistry> registry,
+      std::optional<std::pair<nanobind::object, nanobind::object>> node_data,
+      nanobind::iterable children);
+
   static PyType_Slot slots_[];
 
  private:
diff --git a/jaxlib/pytree_test.py b/jaxlib/pytree_test.py
@@ -37,6 +37,7 @@ def __init__(self, field0, field1):
   def to_iterable(self):
     return [self.field0, self.field1], (None,)
 
+
 def from_iterable(state, values):
   del state
   return ExampleType2(field0=values[0], field1=values[1])
@@ -79,11 +80,38 @@ def testRegisteredType(self):
     with self.assertRaises(ValueError):
       self.roundtrip_proto({"a": ExampleType2(field0=o, field1=o)})
 
+  def roundtrip_node_data(self, example):
+    original = registry.flatten(example)[1]
+    restored = pytree.PyTreeDef.from_node_data_and_children(
+        registry, original.node_data(), original.children()
+    )
+    self.assertEqual(restored, original)
+
+  def testRoundtripNodeData(self):
+    o = object()
+    self.roundtrip_node_data([o, o, o])
+    self.roundtrip_node_data((o, o, o))
+    self.roundtrip_node_data({"a": o, "b": o})
+    self.roundtrip_node_data({22: o, 88: o})
+    self.roundtrip_node_data(None)
+    self.roundtrip_node_data(o)
+    self.roundtrip_node_data(ExampleType(field0=o, field1=o))
+    self.roundtrip_node_data(ExampleType2(field0=o, field1=o))
+
   def testCompose(self):
     x = registry.flatten(0)[1]
     y = registry.flatten((0, 0))[1]
     self.assertEqual((x.compose(y)).num_leaves, 2)
 
+  def testDataclassMakeFromNodeData(self):
+    c = Custom(1, "a")
+    c_leafs, c_tree = registry.flatten(c)
+    c_tree2 = pytree.PyTreeDef.from_node_data_and_children(
+        registry, c_tree.node_data(), c_tree.children()
+    )
+    self.assertEqual(c_tree2.unflatten(c_leafs), c)
+    self.assertEqual(str(c_tree2), str(c_tree))
+
   def testTpTraverse(self):
     self.assertContainsSubset(
         [
