Adds support for symbolic shapes (and their constraints) to OEX V1's OBM path

PiperOrigin-RevId: 738427406

Author: qdhack

export/orbax/export/constants.py

@@ -80,6 +80,9 @@ class OrbaxNativeSerializationType(enum.Enum):
 # attribute on the OrbaxModule.
 WEIGHTS_NAME = 'weights_name'
 
+# Jax2obm_kwargs key for input polymorphic constraints.
+POLYMORPHIC_CONSTRAINTS = 'polymorphic_constraints'
+
 # 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/jax_module.py

@@ -67,8 +67,7 @@ def __init__(
       input_polymorphic_shape: the polymorhpic shape for the inputs of
         ``apply_fn``. If ``apply_fn`` is a mapping, ``input_polymorphic_shape``
         must be a mapping of method key to the input polymorphic shape for the
-        method. Currently input_polymorphic_shape is only relevant for TF
-        SavedModel export.
+        method.
       jax2tf_kwargs: options passed to jax2tf. ``polymorphic_shape`` is inferred
         from ``input_polymorphic_shape`` and should not be set.
         ``with_gradient``, if set, should be consistent with the ``trainable``
@@ -104,6 +103,7 @@ def __init__(
       self._export_module = obm_module.ObmModule(
           params=params,
           apply_fn=apply_fn,
+          input_polymorphic_shape=input_polymorphic_shape,
           jax2obm_kwargs=jax2obm_kwargs,
       )
     else:

export/orbax/export/modules/obm_module.py

@@ -17,16 +17,17 @@
 from collections.abc import Callable, Mapping, Sequence
 import copy
 import logging
-from typing import Any, Optional, Union
+from typing import Any, Optional, Tuple, Union
 
 import jax
 from orbax.export import constants
 from orbax.export import typing as orbax_export_typing
+from orbax.export import utils
 from orbax.export.modules import orbax_module_base
 from orbax.export.typing import PyTree
-from orbax.export import utils
 import tensorflow as tf
 
+
 ApplyFn = orbax_export_typing.ApplyFn
 
 
@@ -37,13 +38,16 @@ def __init__(
       self,
       params: PyTree,
       apply_fn: Union[ApplyFn, Mapping[str, ApplyFn]],
+      *,
+      input_polymorphic_shape: Any = None,
       jax2obm_kwargs: Union[Mapping[str, Any], None] = None,
   ):
     """Data container for Orbax Model export.
 
     Args:
       params: The model parameter specs (e.g. `jax.ShapeDtypeStruct`s).
       apply_fn: The apply_fn for the model.
+      input_polymorphic_shape: polymorhpic shape for the inputs of `apply_fn`.
       jax2obm_kwargs: A dictionary of kwargs to pass to the jax2obm conversion
         library. Accepted arguments to jax2obm_kwargs are
         'native_serialization_platforms', 'flatten_signature', 'weights_name'and
@@ -54,16 +58,29 @@ def __init__(
     if not jax2obm_kwargs:
       jax2obm_kwargs = {}
 
-    self._apply_fn_map = self._normalize_apply_fn_map(
-        self._normalize_apply_fn_map(apply_fn)
+    self._apply_fn_map, self.input_polymorphic_shape_map = (
+        self._normalize_apply_fn_map(apply_fn, input_polymorphic_shape)
     )
 
-    if len(self._apply_fn_map) != 1:
+    if (
+        len(self._apply_fn_map) != 1
+        or len(self.input_polymorphic_shape_map) != 1
+    ):
       raise NotImplementedError(
           'ObmModule: Currently the ObmExport only supports a single method'
-          f' for export. Received: {self._apply_fn_map}'
+          f' for export. Received apply_fn_map: {self._apply_fn_map} and'
+          f' input_polymorphic_shape_map: {self.input_polymorphic_shape_map}'
       )
 
+    # TODO(qidichen): Consider if `self.polymorphic_constraints` should support
+    # a map as well like `input_polymorphic_shape`. For now we assume there is
+    # only one entry in the `apply_fn` or `input_polymorphic_shape`.
+    self.polymorphic_constraints = jax2obm_kwargs.get(
+        constants.POLYMORPHIC_CONSTRAINTS, None
+    )
+    if self.polymorphic_constraints is None:
+      self.polymorphic_constraints = ()
+
     self._native_serialization_platforms = utils.get_lowering_platforms(
         jax2obm_kwargs
     )
@@ -86,18 +103,32 @@ def __init__(
     self._maybe_set_orbax_checkpoint_path(jax2obm_kwargs)
 
   def _normalize_apply_fn_map(
-      self, apply_fn: Union[ApplyFn, Mapping[str, ApplyFn]]
-  ) -> Mapping[str, ApplyFn]:
+      self,
+      apply_fn: Union[ApplyFn, Mapping[str, ApplyFn]],
+      input_polymorphic_shape: Union[PyTree, Mapping[str, PyTree], None],
+  ) -> Tuple[Mapping[str, ApplyFn], Mapping[str, Union[PyTree, None]]]:
     if callable(apply_fn):
       apply_fn_map = {constants.DEFAULT_METHOD_KEY: apply_fn}
+      input_polymorphic_shape_map = {
+          constants.DEFAULT_METHOD_KEY: input_polymorphic_shape
+      }
     elif len(apply_fn) > 1:
       raise NotImplementedError(
           'ObmModule: Currently the ObmExport only supports a single method'
           f' per module. Received: {apply_fn}'
       )
     else:
       apply_fn_map = apply_fn
-    return apply_fn_map
+      if input_polymorphic_shape is None:
+        input_polymorphic_shape_map = {constants.DEFAULT_METHOD_KEY: None}
+      elif not isinstance(input_polymorphic_shape, Mapping):
+        raise TypeError(
+            'When apply_fn is a mapping, input_polymorphic_shape must also be a'
+            ' mapping.'
+        )
+      else:
+        input_polymorphic_shape_map = input_polymorphic_shape
+    return apply_fn_map, input_polymorphic_shape_map
 
   def _maybe_set_orbax_checkpoint_path(self, jax2obm_kwargs):
     if constants.CHECKPOINT_PATH not in jax2obm_kwargs:

export/orbax/export/obm_export.py

@@ -19,6 +19,7 @@
 
 from absl import logging
 import jax
+from jax import export as jax_export
 from orbax.export import constants
 from orbax.export import export_base
 from orbax.export import jax_module

export/orbax/export/obm_export_test.py

@@ -14,7 +14,6 @@
 
 import contextlib
 import os
-from typing import cast
 
 from absl.testing import absltest
 from absl.testing import parameterized