Merge pull request jax-ml#24342 from gnecula:export_custom_types

PiperOrigin-RevId: 688093192

Author: Google-ML-Automation

jax/_src/export/_export.py

@@ -17,13 +17,15 @@
 
 from __future__ import annotations
 
+import collections
 from collections.abc import Callable, Sequence
 import copy
 import dataclasses
 import functools
 import itertools
+import json
 import re
-from typing import Any, Union, cast
+from typing import Any, Protocol, TypeVar, Union, cast
 import warnings
 
 from absl import logging
@@ -344,6 +346,189 @@ def deserialize(blob: bytearray) -> Exported:
   return deserialize(blob)
 
 
+T = TypeVar("T")
+PyTreeAuxData = Any  # alias for tree_util._AuxData
+
+
+class _SerializeAuxData(Protocol):
+  def __call__(self, aux_data: PyTreeAuxData) -> bytes:
+    """Serializes the PyTree node AuxData.
+
+    The AuxData is returned by the `flatten_func` registered by
+    `tree_util.register_pytree_node`).
+    """
+
+
+class _DeserializeAuxData(Protocol):
+  def __call__(self, serialized_aux_data: bytes) -> PyTreeAuxData:
+    """Deserializes the PyTree node AuxData.
+
+    The result will be passed to `_BuildFromChildren`.
+    """
+
+
+class _BuildFromChildren(Protocol):
+  def __call__(self, aux_data: PyTreeAuxData, children: Sequence[Any]) -> Any:
+    """Materializes a T given a deserialized AuxData and children.
+
+    This is similar in scope with the `unflatten_func`.
+    """
+
+
+serialization_registry: dict[type, tuple[str, _SerializeAuxData]] = {}
+
+
+deserialization_registry: dict[
+  str,
+  tuple[type, _DeserializeAuxData, _BuildFromChildren]] = {}
+
+
+def _is_namedtuple(nodetype: type) -> bool:
+  return (issubclass(nodetype, tuple) and
+          hasattr(nodetype, "_fields") and
+          isinstance(nodetype._fields, Sequence) and
+          all(isinstance(f, str) for f in nodetype._fields))
+
+def register_pytree_node_serialization(
+    nodetype: type[T],
+    *,
+    serialized_name: str,
+    serialize_auxdata: _SerializeAuxData,
+    deserialize_auxdata: _DeserializeAuxData,
+    from_children: _BuildFromChildren | None = None
+) -> type[T]:
+  """Registers a custom PyTree node for serialization and deserialization.
+
+  You must use this function before you can serialize and deserialize PyTree
+  nodes for the types not supported natively. We serialize PyTree nodes for
+  the `in_tree` and `out_tree` fields of `Exported`, which are part of the
+  exported function's calling convention.
+
+  This function must be called after calling
+  `jax.tree_util.register_pytree_node` (except for `collections.namedtuple`,
+  which do not require a call to `register_pytree_node`).
+
+  Args:
+    nodetype: the type whose PyTree nodes we want to serialize. It is an
+      error to attempt to register multiple serializations for a `nodetype`.
+    serialized_name: a string that will be present in the serialization and
+      will be used to look up the registration during deserialization. It is an
+      error to attempt to register multiple serializations for a
+      `serialized_name`.
+    serialize_auxdata: serialize the PyTree auxdata (returned by the
+      `flatten_func` argument to `jax.tree_util.register_pytree_node`.).
+    deserialize_auxdata: deserialize the auxdata that was serialized by the
+      `serialize_auxdata`.
+    from_children: if present, this is a function that takes that result of
+      `deserialize_auxdata` along with some children and creates an instance
+      of `nodetype`. This is similar to the `unflatten_func` passed to
+      `jax.tree_util.register_pytree_node`. If not present, we look up
+      and use the `unflatten_func`. This is needed for `collections.namedtuple`,
+      which does not have a `register_pytree_node`, but it can be useful to
+      override that function. Note that the result of `from_children` is
+      only used with `jax.tree_util.tree_structure` to construct a proper
+      PyTree node, it is not used to construct the outputs of the serialized
+      function.
+
+  Returns:
+    the same type passed as `nodetype`, so that this function can
+    be used as a class decorator.
+  """
+  if nodetype in serialization_registry:
+    raise ValueError(
+        f"Duplicate serialization registration for type `{nodetype}`. "
+        "Previous registration was with serialized_name "
+        f"`{serialization_registry[nodetype][0]}`.")
+  if serialized_name in deserialization_registry:
+    raise ValueError(
+        "Duplicate serialization registration for "
+        f"serialized_name `{serialized_name}`. "
+        "Previous registration was for type "
+        f"`{deserialization_registry[serialized_name][0]}`.")
+  if from_children is None:
+    if nodetype not in tree_util._registry:
+      raise ValueError(
+          f"If `from_children` is not present, you must call first"
+          f"`jax.tree_util.register_pytree_node` for `{nodetype}`")
+    from_children = tree_util._registry[nodetype].from_iter
+
+  serialization_registry[nodetype] = (
+      serialized_name, serialize_auxdata)
+  deserialization_registry[serialized_name] = (
+      nodetype, deserialize_auxdata, from_children)
+  return nodetype
+
+
+def register_namedtuple_serialization(
+    nodetype: type[T],
+    *,
+    serialized_name: str) -> type[T]:
+  """Registers a namedtuple for serialization and deserialization.
+
+  JAX has native PyTree support for `collections.namedtuple`, and does not
+  require a call to `jax.tree_util.register_pytree_node`. However, if you
+  want to serialize functions that have inputs of outputs of a
+  namedtuple type, you must register that type for serialization.
+
+  Args:
+    nodetype: the type whose PyTree nodes we want to serialize. It is an
+      error to attempt to register multiple serializations for a `nodetype`.
+      On deserialization, this type must have the same set of keys that
+      were present during serialization.
+    serialized_name: a string that will be present in the serialization and
+      will be used to look up the registration during deserialization. It is an
+      error to attempt to register multiple serializations for
+      a `serialized_name`.
+
+  Returns:
+    the same type passed as `nodetype`, so that this function can
+    be used as a class decorator.
+"""
+  if not _is_namedtuple(nodetype):
+    raise ValueError("Use `jax.export.register_pytree_node_serialization` for "
+                     "types other than `collections.namedtuple`.")
+
+  def serialize_auxdata(aux_data: PyTreeAuxData) -> bytes:
+    # Store the serialized keys in the serialized auxdata
+    del aux_data
+    return json.dumps(nodetype._fields).encode("utf-8")
+
+  def deserialize_auxdata(serialized_aux_data: bytes) -> PyTreeAuxData:
+    return json.loads(serialized_aux_data.decode("utf-8"))
+
+  def from_children(aux_data: PyTreeAuxData, children: Sequence[Any]) -> Any:
+    # Use our own "from_children" because namedtuples do not have a pytree
+    # registration.
+    ser_keys = cast(Sequence[str], aux_data)
+    assert len(ser_keys) == len(children)
+    return nodetype(** dict(zip(ser_keys, children)))
+
+  return register_pytree_node_serialization(
+      nodetype,
+      serialized_name=serialized_name,
+      serialize_auxdata=serialize_auxdata,
+      deserialize_auxdata=deserialize_auxdata,
+      from_children=from_children)
+
+
+# collections.OrderedDict is registered as a pytree node with auxdata being
+# `tuple(x.keys())`.
+def _serialize_ordereddict_keys(keys):
+  if isinstance(keys, Sequence) and all(isinstance(k, str) for k in keys):
+    return json.dumps(keys).encode("utf-8")
+  else:
+    raise NotImplementedError(
+        "Serialization of collections.OrderedDict is supported only when the "
+        f"keys are strings. Found keys: {keys}.")
+
+
+register_pytree_node_serialization(
+    collections.OrderedDict,
+    serialized_name="collections.OrderedDict",
+    serialize_auxdata=_serialize_ordereddict_keys,
+    deserialize_auxdata=lambda b: json.loads(b.decode("utf-8")))
+
+
 def default_export_platform() -> str:
   """Retrieves the default export platform.
 
@@ -404,9 +589,10 @@ def export_back_compat(
     disabled_checks: the safety checks to disable. See docstring
         of `DisabledSafetyCheck`.
 
-  Returns: a function that takes args and kwargs pytrees of jax.ShapeDtypeStruct,
-      or values with `.shape` and `.dtype` attributes, and returns an
-      `Exported`.
+  Returns:
+    a function that takes args and kwargs pytrees of jax.ShapeDtypeStruct,
+    or values with `.shape` and `.dtype` attributes, and returns an
+    `Exported`.
 
   Usage:
 
@@ -480,9 +666,10 @@ def export(
     disabled_checks: the safety checks to disable. See documentation for
         of `jax.export.DisabledSafetyCheck`.
 
-  Returns: a function that takes args and kwargs pytrees of {class}`jax.ShapeDtypeStruct`,
-      or values with `.shape` and `.dtype` attributes, and returns an
-      `Exported`.
+  Returns:
+    a function that takes args and kwargs pytrees of {class}`jax.ShapeDtypeStruct`,
+    or values with `.shape` and `.dtype` attributes, and returns an
+    `Exported`.
 
   Usage:
 

jax/_src/export/serialization.fbs

@@ -28,12 +28,15 @@ enum PyTreeDefKind: byte {
   tuple = 2,
   list = 3,
   dict = 4,
+  custom = 5,
 }
 
 table PyTreeDef {
   kind: PyTreeDefKind;
   children: [PyTreeDef];
-  children_names: [string];  // only for "dict"
+  children_names: [string];  // only for "kind==dict"
+  custom_name: string;  // only for "kind==custom"
+  custom_auxdata: [byte];  // only for "kind==custom"
 }
 
 enum AbstractValueKind: byte {

jax/_src/export/serialization.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import types
 from collections.abc import Callable, Sequence
 from functools import partial
 from typing import TypeVar
@@ -45,6 +46,8 @@
 # even if the change is backwards compatible.
 # Version 1, Nov 2023, first version.
 # Version 2, Dec 16th, 2023, adds the f0 dtype.
+# Version 3, October 16th, 2024, adds serialization for namedtuple and custom types
+#   This version is backwards compatible with Version 2.
 _SERIALIZATION_VERSION = 2
 
 def serialize(exp: _export.Exported, vjp_order: int = 0) -> bytearray:
@@ -152,21 +155,21 @@ def _serialize_array(
 
 def _deserialize_exported(exp: ser_flatbuf.Exported) -> _export.Exported:
   serialization_version = exp.SerializationVersion()
-  if serialization_version != _SERIALIZATION_VERSION:
+  if serialization_version not in [2, 3]:
     raise NotImplementedError(
         f"deserialize unsupported version {serialization_version}"
     )
 
   fun_name = exp.FunctionName().decode("utf-8")
-  _, in_tree = tree_util.tree_flatten(
+  in_tree = tree_util.tree_structure(
       _deserialize_pytreedef_to_pytree(exp.InTree())
   )
   scope = shape_poly.SymbolicScope(())  # TODO: serialize the constraints
   deser_aval = partial(_deserialize_aval, scope=scope)
   in_avals = _deserialize_tuple(
       exp.InAvalsLength, exp.InAvals, deser_aval
   )
-  _, out_tree = tree_util.tree_flatten(
+  out_tree = tree_util.tree_structure(
       _deserialize_pytreedef_to_pytree(exp.OutTree())
   )
   out_avals = _deserialize_tuple(
@@ -246,30 +249,51 @@ def _serialize_pytreedef(
     children_vector_offset = _serialize_array(
         builder, _serialize_pytreedef, children
     )
+  custom_name = None
+  custom_auxdata = None
+  node_type = node_data and node_data[0]
 
   if node_data is None:  # leaf
     kind = ser_flatbuf.PyTreeDefKind.leaf
-  elif node_data[0] is type(None):
+  elif node_type is types.NoneType:
     kind = ser_flatbuf.PyTreeDefKind.none
-  elif node_data[0] is tuple:
+  elif node_type is tuple:
     kind = ser_flatbuf.PyTreeDefKind.tuple
-  elif node_data[0] is list:
+  elif node_type is list:
     kind = ser_flatbuf.PyTreeDefKind.list
-  elif node_data[0] is dict:
+  elif node_type is dict:
     kind = ser_flatbuf.PyTreeDefKind.dict
     assert len(node_data[1]) == len(children)
     children_names_vector_offset = _serialize_array(
         builder, lambda b, s: b.CreateString(s), node_data[1]
     )
+  elif node_type in _export.serialization_registry:
+    kind = ser_flatbuf.PyTreeDefKind.custom
+    serialized_name, serialize_auxdata = _export.serialization_registry[node_type]
+    custom_name = builder.CreateString(serialized_name)
+    serialized_auxdata = serialize_auxdata(node_data[1])
+    if not isinstance(serialized_auxdata, (bytes, bytearray)):
+      raise ValueError(
+          "The custom serialization function for `node_type` must "
+          f"return a `bytes` object. It returned a {type(serialized_auxdata)}.")
+    custom_auxdata = builder.CreateByteVector(serialized_auxdata)
   else:
-    raise NotImplementedError(f"serializing PyTreeDef {node_data}")
+    raise ValueError(
+        "Cannot serialize PyTreeDef containing an "
+        f"unregistered type `{node_type}`. "
+        "Use `export.register_pytree_node_serialization` or "
+        "`export.register_namedtuple_serialization`.")
 
   ser_flatbuf.PyTreeDefStart(builder)
   ser_flatbuf.PyTreeDefAddKind(builder, kind)
   if children_vector_offset:
     ser_flatbuf.PyTreeDefAddChildren(builder, children_vector_offset)
   if children_names_vector_offset:
     ser_flatbuf.PyTreeDefAddChildrenNames(builder, children_names_vector_offset)
+  if custom_name is not None:
+    ser_flatbuf.PyTreeDefAddCustomName(builder, custom_name)
+  if custom_auxdata is not None:
+    ser_flatbuf.PyTreeDefAddCustomAuxdata(builder, custom_auxdata)
   return ser_flatbuf.PyTreeDefEnd(builder)
 
 
@@ -294,6 +318,17 @@ def _deserialize_pytreedef_to_pytree(p: ser_flatbuf.PyTreeDef):
     assert p.ChildrenNamesLength() == nr_children
     keys = [p.ChildrenNames(i).decode("utf-8") for i in range(nr_children)]
     return dict(zip(keys, children))
+  elif kind == ser_flatbuf.PyTreeDefKind.custom:
+    serialized_name = p.CustomName().decode("utf-8")
+    if serialized_name not in _export.deserialization_registry:
+      raise ValueError(
+          "Cannot deserialize a PyTreeDef containing an "
+          f"unregistered type `{serialized_name}`. "
+          "Use `export.register_pytree_node_serialization` or "
+          "`export.register_namedtuple_serialization`.")
+    nodetype, deserialize_auxdata, from_iter = _export.deserialization_registry[serialized_name]
+    auxdata = deserialize_auxdata(p.CustomAuxdataAsNumpy().tobytes())
+    return from_iter(auxdata, children)
   else:
     assert False, kind