Add save spec methods to TF layer

keras_core/backend/tensorflow/layer.py

@@ -1,11 +1,59 @@
 import tensorflow as tf
-
+from keras_core.backend.tensorflow import utils as tf_utils
 
 class TFLayer(tf.__internal__.tracking.AutoTrackable):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Export-related attributes
+        self._saved_model_inputs_spec = None
+        self._saved_model_arg_spec = None
+
     def _post_build(self):
         """Can be overriden to perform post-build actions."""
         pass
 
+    @tf.__internal__.tracking.no_automatic_dependency_tracking
+    def _set_save_spec(self, inputs, args=None, kwargs=None):
+        """Defines the save spec so that serialization can trace layer calls.
+
+        The TensorSpecs of the call function `inputs`, `args`, and `kwargs` are
+        saved into a tuple of `([inputs] + args, kwargs)`.
+
+        Args:
+          inputs: possibly nested inputs passed into the call function.
+          args: a list of positional arguments passed into call.
+          kwargs: a dictionary of keyword arguments passed into call.
+        """
+        if self._saved_model_inputs_spec is not None:
+            return  # Already set.
+
+        inputs_spec = tf.nest.map_structure(tf_utils.get_tensor_spec, inputs)
+        args_spec = tf.nest.map_structure(tf_utils.get_tensor_spec, args or [])
+        kwargs_spec = {}
+        # Filter out non-tensor arguments from kwargs.
+        for key, kwarg in kwargs.items():
+            flat_kwarg = tf.nest.flatten(kwarg)
+            flat_specs = [tf_utils.get_tensor_spec(x) for x in flat_kwarg]
+            if any(s is None for s in flat_specs):
+                continue
+            kwargs_spec[key] = tf.nest.pack_sequence_as(kwarg, flat_specs)
+
+        self._saved_model_inputs_spec = inputs_spec
+        self._saved_model_arg_spec = (
+            [inputs_spec] + list(args_spec),
+            kwargs_spec,
+        )
+
+    def _get_save_spec(self, dynamic_batch=True, inputs_only=True):
+        if self._saved_model_inputs_spec is None:
+            return None
+
+        spec = tf.nest.map_structure(
+            lambda t: tf_utils.get_tensor_spec(t, dynamic_batch=dynamic_batch),
+            self._saved_model_arg_spec,
+        )
+        return spec[0][0] if inputs_only else spec
+
     def _trackable_children(self, save_type="checkpoint", **kwargs):
         if save_type == "savedmodel":
             # SavedModel needs to ignore the execution functions.

keras_core/backend/tensorflow/utils.py

@@ -0,0 +1,26 @@
+import tensorflow as tf
+
+def get_tensor_spec(t, dynamic_batch=False, name=None):
+    """Returns a `TensorSpec` given a single `Tensor` or `TensorSpec`."""
+    if isinstance(t, tf.TypeSpec):
+        spec = t
+    elif isinstance(tensor, tf.__internal__.CompositeTensor):
+        # Check for ExtensionTypes
+        spec = t._type_spec
+    elif hasattr(t, "shape") and hasattr(t, "dtype"):
+        spec = tf.TensorSpec(shape=t.shape, dtype=t.dtype, name=name)
+    else:
+        return None  # Allow non-Tensors to pass through.
+
+    if not dynamic_batch:
+        return spec
+
+    shape = spec.shape
+    if shape.rank is None or shape.rank == 0:
+        return spec
+
+    shape_list = shape.as_list()
+    shape_list[0] = None
+    shape = tf.TensorShape(shape_list)
+    spec._shape = shape
+    return spec

keras_core/export/export_lib_test.py

@@ -6,6 +6,7 @@
 import pytest
 from absl.testing import parameterized
 
+from keras_core import backend
 from keras_core import testing
 from keras_core import layers
 from keras_core import models
@@ -23,6 +24,10 @@ def get_model():
     return model
 
 
+@pytest.mark.skipif(
+    backend.backend() != "tensorflow",
+    reason="Export only currently supports the TF backend.",
+)
 class ExportArchiveTest(testing.TestCase, parameterized.TestCase):
     def test_standard_model_export(self):
         temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
@@ -452,6 +457,10 @@ def test_model_export_method(self):
         )
 
 
+@pytest.mark.skipif(
+    backend.backend() != "tensorflow",
+    reason="Export only currently supports the TF backend.",
+)
 class TestReloadedLayer(tf.test.TestCase, parameterized.TestCase):
     def test_reloading_export_archive(self):
         temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")