Adds a function to check the inputs have default values in orbax savedmodel.

checkpoint/CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add `assert_tensor_spec_with_default` for input signature testing
 - Add support for loading SafeTensors checkpoints
 - #v1 Add `is_orbax_checkpoint()` method for validation checks
 

export/orbax/export/utils.py

@@ -89,6 +89,16 @@ def __post_init__(self):
 ]
 
 
+def assert_tensor_spec_with_default(
+    input_signature: PyTree,
+) -> PyTree:
+  """Asserts that the input signature is a TensorSpecWithDefault."""
+  def check_fn(x):
+    assert isinstance(x, TensorSpecWithDefault), f'x: {x}'
+    return x
+  return jax.tree_util.tree_map(check_fn, input_signature)
+
+
 def remove_signature_defaults(input_signature: PyTree) -> PyTree:
   """Removes TensorSpecWithDefault from an input_signature."""
 

export/orbax/export/utils_test.py

@@ -65,6 +65,24 @@ def test_missing_default(self):
     ):
       utils.with_default_args(lambda x: x[0] + x[1], input_signature)
 
+  def test_assert_tensor_spec_with_default(self):
+    input_signature = [
+        TensorSpecWithDefault(
+            tf.TensorSpec([None], tf.int32),
+            np.asarray([1, 2]),
+        )
+    ]
+    utils.assert_tensor_spec_with_default(input_signature)
+
+    input_signature_bad_type = [
+        tf.TensorSpec([None], tf.int32),
+    ]
+    with self.assertRaisesRegex(
+        AssertionError,
+        'x: TensorSpec',
+    ):
+      utils.assert_tensor_spec_with_default(input_signature_bad_type)
+
   def test_with_default_args_nested(self):
     def f(required_arg, optional_args):
       return (