Model Export to liteRT

keras/src/backend/tensorflow/layer.py

@@ -13,7 +13,6 @@ def __init__(self, *args, **kwargs):
         self._saved_model_arg_spec = None
         self._tracked = []
 
-    @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.
 
@@ -45,6 +44,7 @@ def _set_save_spec(self, inputs, args=None, kwargs=None):
             kwargs_spec,
         )
 
+    @tf.__internal__.tracking.no_automatic_dependency_tracking
     def _trackable_children(self, save_type="checkpoint", **kwargs):
         if save_type == "savedmodel":
             # SavedModel needs to ignore the execution functions.
@@ -62,17 +62,55 @@ def _trackable_children(self, save_type="checkpoint", **kwargs):
             self.test_function = test_function
             self.predict_function = predict_function
 
-            for tracked_attr in self._tracked:
-                tracked_item = getattr(self, tracked_attr)
-                if isinstance(tracked_item, tracking.TrackedList):
-                    children[tracked_attr] = list(tracked_item)
-                if isinstance(tracked_item, tracking.TrackedDict):
-                    children[tracked_attr] = dict(tracked_item)
-                if isinstance(tracked_item, tracking.TrackedSet):
-                    children[tracked_attr] = list(tracked_item)
+            # Convert Keras tracked collections to plain Python structures
+            # without creating TensorFlow trackable dependencies
+            self._convert_tracked_collections(children)
 
         return children
 
+    def _convert_tracked_collections(self, children):
+        """Convert TrackedList/Dict/Set to plain Python structures."""
+        for tracked_attr in self._tracked:
+            tracked_item = getattr(self, tracked_attr)
+            if isinstance(tracked_item, tracking.TrackedList):
+                children[tracked_attr] = list(tracked_item)
+            if isinstance(tracked_item, tracking.TrackedDict):
+                children[tracked_attr] = dict(tracked_item)
+            if isinstance(tracked_item, tracking.TrackedSet):
+                children[tracked_attr] = list(tracked_item)
+
+    def _get_save_spec(self, dynamic_batch=True):
+        """Compatibility shim for TensorFlow saving utilities.
+
+        TensorFlow's SavedModel / TFLite export paths (e.g.,
+        tf.lite.TFLiteConverter.from_keras_model) expect a `_get_save_spec`
+        method on models. This method generates TensorSpec objects
+        describing the model's input signature.
+
+        Args:
+            dynamic_batch: whether to set the batch dimension to `None`.
+
+        Returns:
+            A TensorSpec, list or dict mirroring the model inputs, or
+            `None` when specs cannot be inferred.
+        """
+        # Prefer the base implementation if available
+        try:
+            return super()._get_save_spec(dynamic_batch)
+        except AttributeError:
+            # Lazy import to avoid circular dependency
+            from keras.src.export.export_utils import make_tf_tensor_spec
+
+            # Fall back to building specs from `self.inputs`
+            inputs = getattr(self, "inputs", None)
+            if inputs is None:
+                return None
+
+            return tree.map_structure(
+                lambda x: make_tf_tensor_spec(x, dynamic_batch=dynamic_batch),
+                inputs,
+            )
+
     @property
     def _default_save_signature(self):
         """For SavedModel support: returns the default serving signature."""

keras/src/export/__init__.py

@@ -1,3 +1,5 @@
+from keras.src.export.litert import LiteRTExporter
+from keras.src.export.litert import export_litert
 from keras.src.export.onnx import export_onnx
 from keras.src.export.openvino import export_openvino
 from keras.src.export.saved_model import ExportArchive

keras/src/export/export_utils.py

@@ -7,6 +7,14 @@
 
 
 def get_input_signature(model):
+    """Get input signature for model export.
+
+    Args:
+        model: A Keras Model instance.
+
+    Returns:
+        Input signature suitable for model export (always a tuple or list).
+    """
     if not isinstance(model, models.Model):
         raise TypeError(
             "The model must be a `keras.Model`. "
@@ -17,13 +25,20 @@ def get_input_signature(model):
             "The model provided has not yet been built. It must be built "
             "before export."
         )
+
     if isinstance(model, models.Functional):
+        # Functional models expect a single positional argument `inputs`
+        # containing the full nested input structure. We keep the
+        # original behavior of returning a single-element list that
+        # wraps the mapped structure so that downstream exporters
+        # build a tf.function with one positional argument.
         input_signature = [
             tree.map_structure(make_input_spec, model._inputs_struct)
         ]
     elif isinstance(model, models.Sequential):
         input_signature = tree.map_structure(make_input_spec, model.inputs)
     else:
+        # Subclassed models: rely on recorded shapes from the first call.
         input_signature = _infer_input_signature_from_model(model)
         if not input_signature or not model._called:
             raise ValueError(
@@ -60,6 +75,7 @@ def _make_input_spec(structure):
                 f"Unsupported type {type(structure)} for {structure}"
             )
 
+    # Always return a flat list preserving the order of shapes_dict values
     return [_make_input_spec(value) for value in shapes_dict.values()]
 
 
@@ -86,13 +102,33 @@ def make_input_spec(x):
     return input_spec
 
 
-def make_tf_tensor_spec(x):
+def make_tf_tensor_spec(x, dynamic_batch=True):
+    """Create a TensorSpec from various input types.
+
+    Args:
+        x: Input to convert (tf.TensorSpec, KerasTensor, or backend tensor).
+        dynamic_batch: If True, set the batch dimension to None.
+
+    Returns:
+        A tf.TensorSpec instance.
+    """
     if isinstance(x, tf.TensorSpec):
         tensor_spec = x
+        # Adjust batch dimension if needed
+        if dynamic_batch and len(tensor_spec.shape) > 0:
+            shape = list(tensor_spec.shape)
+            shape[0] = None
+            tensor_spec = tf.TensorSpec(
+                tuple(shape), dtype=tensor_spec.dtype, name=tensor_spec.name
+            )
     else:
         input_spec = make_input_spec(x)
+        shape = input_spec.shape
+        # Adjust batch dimension if needed and shape is not None
+        if dynamic_batch and shape is not None and len(shape) > 0:
+            shape = tuple(None if i == 0 else s for i, s in enumerate(shape))
         tensor_spec = tf.TensorSpec(
-            input_spec.shape, dtype=input_spec.dtype, name=input_spec.name
+            shape, dtype=input_spec.dtype, name=input_spec.name
         )
     return tensor_spec