Add prune_tree utility and integrate it into Orbax Model export.

This change introduces a `prune_tree` function to `tree_util.py` which removes tree leaves not matching a specified type. The Orbax Model export pipeline is updated with a new option, `prune_custom_pytree_nodes`. When this option is enabled, custom PyTree nodes that are not `ShloTensorSpec` instances are pruned (replaced with `None`) instead of causing an error during signature conversion. This allows for more flexible handling of PyTrees containing non-tensor metadata.

PiperOrigin-RevId: 810549387

Author: jerryxyj

export/orbax/export/constants.py

@@ -62,6 +62,9 @@ class ExportModelType(enum.Enum):
 # Jax2obm_kwargs key for input polymorphic constraints.
 POLYMORPHIC_CONSTRAINTS = 'polymorphic_constraints'
 
+# Jax2obm_kwargs key for pruning custom pytree nodes.
+PRUNE_CUSTOM_PYTREE_NODES = 'prune_custom_pytree_nodes'
+
 # Default weights name to use if a checkpoint path is provided but a weights_
 # name kwarg was not provided in the jax2obm_kwargs.
 DEFAULT_WEIGHTS_NAME = 'weights'

export/orbax/export/modules/obm_module.py

@@ -104,6 +104,9 @@ def __init__(
     self._checkpoint_path: str = None
     # Set the Orbax checkpoint path if provided in the jax2obm_kwargs.
     self._maybe_set_orbax_checkpoint_path(jax2obm_kwargs)
+    self._prune_custom_pytree_nodes = jax2obm_kwargs.get(
+        constants.PRUNE_CUSTOM_PYTREE_NODES, False
+    )
 
   def _normalize_apply_fn_map(
       self,

export/orbax/export/obm_export_test.py

@@ -15,6 +15,7 @@
 from collections.abc import Mapping, Sequence
 import contextlib
 import os
+from typing import Any, Callable
 
 from absl.testing import absltest
 from absl.testing import parameterized

model/orbax/experimental/model/core/python/tree_util.py

@@ -142,3 +142,19 @@ def unflatten(tree: Tree[Any], leaves: Iterable[T6]) -> Tree[T6]:
   except StopIteration:
     return result
   raise ValueError("After unflattening, there are still leaves left.")
+
+
+def prune_tree(tree: Tree[Any], wanted_node_type: Any) -> Tree[Any]:
+  """Prunes the tree by removing unwanted leaves."""
+  if isinstance(tree, (tuple, list)):
+    return tuple_or_list_constructor(tree)(
+        prune_tree(x, wanted_node_type) for x in tree
+    )
+  elif isinstance(tree, dict):
+    tree: Dict[str, Tree[Any]]
+    return {k: prune_tree(v, wanted_node_type) for k, v in tree.items()}
+  elif isinstance(tree, wanted_node_type):
+    return tree
+  else:
+    # If the node type is None or not wanted, we return None.
+    return None

model/orbax/experimental/model/core/python/tree_util_test.py

@@ -77,6 +77,28 @@ def assert_int(x: Any) -> None:
         lambda: tree_util.assert_tree(assert_int, wrong_tree),
     )
 
+  def test_prune_tree(self):
+    tree = (1, "a", [2, "b", {"c": 3, "d": "e"}], None)
+    pruned_int = tree_util.prune_tree(tree, int)
+    self.assertEqual(
+        pruned_int, (1, None, [2, None, {"c": 3, "d": None}], None)
+    )
+    pruned_str = tree_util.prune_tree(tree, str)
+    self.assertEqual(
+        pruned_str, (None, "a", [None, "b", {"c": None, "d": "e"}], None)
+    )
+    pruned_int_str = tree_util.prune_tree(tree, (int, str))
+    self.assertEqual(pruned_int_str, tree)
+    pruned_empty = tree_util.prune_tree(tree, ())
+    self.assertEqual(
+        pruned_empty, (None, None, [None, None, {"c": None, "d": None}], None)
+    )
+    pruned_not_present = tree_util.prune_tree(tree, float)
+    self.assertEqual(
+        pruned_not_present,
+        (None, None, [None, None, {"c": None, "d": None}], None),
+    )
+
 
 if __name__ == "__main__":
   absltest.main()

model/orbax/experimental/model/jax2obm/jax_specific_info.py

@@ -311,6 +311,7 @@ def _to_shlo_spec_tree_and_refinement_tuple(
     avals: Sequence[jax.core.AbstractValue],
     shardings: Sequence[Any],
     tree_def: Optional[jax.tree_util.PyTreeDef],
+    prune_custom_pytree_nodes: bool = False,
 ) -> Tuple[
     obm.Tree[obm.ShloTensorSpec], Tuple[ShapeDTypeRefinementPair, ...] | None
 ]:
@@ -327,7 +328,8 @@ def assert_leaf(x: Any) -> None:
         raise ValueError(
             f"Leaf needs to be a ShloTensorSpec, but its type is: {type(x)}"
         )
-
+    if prune_custom_pytree_nodes:
+      jax_tree = obm.tree_util.prune_tree(jax_tree, obm.ShloTensorSpec)
     obm.tree_util.assert_tree(assert_leaf, jax_tree)
     jax_tree: obm.Tree[obm.ShloTensorSpec]
   return jax_tree, refinements

model/orbax/experimental/model/jax2obm/main_lib.py

@@ -41,6 +41,7 @@ def jax_exported_to_shlo_fn(
     exported: jax_export.Exported,
     xla_compile_options_per_platform: (
         obm.manifest_pb2.CompileOptionsProtoMap | None
+    prune_custom_pytree_nodes: bool = False,
 ) -> obm.ShloFunction:
   """Converts a `jax.export.Exported` to an Orbax Model `ShloFunction`."""
 
@@ -57,13 +58,15 @@ def jax_exported_to_shlo_fn(
           exported.in_avals,
           in_shardings_hlo,
           exported.in_tree,
+          prune_custom_pytree_nodes=prune_custom_pytree_nodes,
       )
   )
   shlo_out_sig, jax_out_sig_refinements = (
       jax_specific_info._to_shlo_spec_tree_and_refinement_tuple(
           exported.out_avals,
           out_shardings_hlo,
           exported.out_tree,
+          prune_custom_pytree_nodes=prune_custom_pytree_nodes,
       )
   )
   supplemental_info_ = {}
@@ -107,6 +110,7 @@ def convert(
     native_serialization_disabled_checks: Sequence[
         jax_export.DisabledSafetyCheck
     ] = (),
+    prune_custom_pytree_nodes: bool = False,
 ) -> obm.ShloFunction:
   """Converts a JAX function to an Orbax Model `ShloFunction`.
 
@@ -138,6 +142,8 @@ def convert(
       model artifact, to ensure XLA compilation consistency and reproducibility
       between export time and serving time. Each map entry corresponds to a
       platform type (e.g. TPU, GPU, etc.).
+    prune_custom_pytree_nodes: Optional. True if the custom pytree nodes should
+      be pruned. False by default.
 
   Returns:
     An Orbax Model `ShloFunction`.
@@ -156,6 +162,7 @@ def convert(
   exported = exported_creator(*args_spec, **kwargs_spec)
   return jax_exported_to_shlo_fn(
       exported,
+      prune_custom_pytree_nodes=prune_custom_pytree_nodes,
   )