Model Export to liteRT

keras/src/backend/tensorflow/layer.py

@@ -62,17 +62,73 @@ 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
 
+    @tf.__internal__.tracking.no_automatic_dependency_tracking
+    def _convert_tracked_collections(self, children):
+        """Convert TrackedList/Dict/Set to plain Python structures.
+
+        The decorator prevents TensorFlow from automatically wrapping
+        these conversions in _DictWrapper objects.
+        """
+        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:
+            # Fall back to building specs from `self.inputs`
+            inputs = getattr(self, "inputs", None)
+            if inputs is None:
+                return None
+
+            def _make_spec(t):
+                # t is a tf.Tensor-like object
+                shape = list(t.shape)
+                if dynamic_batch and len(shape) > 0:
+                    shape[0] = None
+                # Convert to tuple for TensorSpec
+                try:
+                    name = getattr(t, "name", None)
+                    return tf.TensorSpec(
+                        shape=tuple(shape), dtype=t.dtype, name=name
+                    )
+                except (ImportError, ModuleNotFoundError):
+                    return None
+
+            # Handle dict/list/single tensor inputs
+            if isinstance(inputs, dict):
+                return {k: _make_spec(v) for k, v in inputs.items()}
+            if isinstance(inputs, (list, tuple)):
+                return [_make_spec(t) for t in inputs]
+            return _make_spec(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()]
 
 

keras/src/export/litert.py

@@ -0,0 +1,255 @@
+from keras.src import layers
+from keras.src import models
+from keras.src.export.export_utils import get_input_signature
+from keras.src.export.export_utils import make_input_spec
+from keras.src.utils import io_utils
+from keras.src.utils.module_utils import tensorflow as tf
+
+
+def export_litert(
+    model,
+    filepath,
+    input_signature=None,
+    **kwargs,
+):
+    """Export the model as a LiteRT artifact for inference.
+
+    Args:
+        model: The Keras model to export.
+        filepath: The path to save the exported artifact.
+        input_signature: Optional input signature specification. If
+            `None`, it will be inferred.
+        **kwargs: Additional keyword arguments passed to the exporter.
+    """
+
+    exporter = LiteRTExporter(
+        model=model,
+        input_signature=input_signature,
+        **kwargs,
+    )
+    exporter.export(filepath)
+    io_utils.print_msg(f"Saved artifact at '{filepath}'.")
+
+
+class LiteRTExporter:
+    """Exporter for the LiteRT (TFLite) format.
+
+    This class handles the conversion of Keras models for LiteRT runtime and
+    generates a `.tflite` model file. For efficient inference on mobile and
+    embedded devices, it creates a single callable signature based on the
+    model's `call()` method.
+    """
+
+    def __init__(
+        self,
+        model,
+        input_signature=None,
+        **kwargs,
+    ):
+        """Initialize the LiteRT exporter.
+
+        Args:
+            model: The Keras model to export
+            input_signature: Input signature specification (e.g., TensorFlow
+                TensorSpec or list of TensorSpec)
+            **kwargs: Additional export parameters
+        """
+        self.model = model
+        self.input_signature = input_signature
+        self.kwargs = kwargs
+
+    def _infer_dict_input_signature(self):
+        """Infer input signature from a model with dict inputs.
+
+        This reads the actual shapes and dtypes from model._inputs_struct.
+
+        Returns:
+            dict or None: Dictionary mapping input names to InputSpec, or None
+        """
+        # Check _inputs_struct first (preserves dict structure)
+        if hasattr(self.model, "_inputs_struct") and isinstance(
+            self.model._inputs_struct, dict
+        ):
+            return {
+                name: make_input_spec(inp)
+                for name, inp in self.model._inputs_struct.items()
+            }
+
+        # Fall back to model.inputs if it's a dict
+        if hasattr(self.model, "inputs") and isinstance(
+            self.model.inputs, dict
+        ):
+            return {
+                name: make_input_spec(inp)
+                for name, inp in self.model.inputs.items()
+            }
+
+        return None
+
+    def export(self, filepath):
+        """Exports the Keras model to a TFLite file.
+
+        Args:
+            filepath: Output path for the exported model
+
+        Returns:
+            Path to exported model
+        """
+        # 1. Resolve / infer input signature
+        if self.input_signature is None:
+            # Try dict-specific inference first (for models with dict inputs)
+            dict_signature = self._infer_dict_input_signature()
+            if dict_signature is not None:
+                self.input_signature = dict_signature
+            else:
+                # Fall back to standard inference
+                self.input_signature = get_input_signature(self.model)
+
+        # 3. Handle dictionary inputs by creating an adapter
+        # Check if we have dict inputs that need adaptation
+        has_dict_inputs = isinstance(self.input_signature, dict)
+
+        if has_dict_inputs:
+            # Create adapter model that converts list to dict
+            adapted_model = self._create_dict_adapter(self.input_signature)
+
+            # Convert dict signature to list for TFLite conversion
+            # The adapter will handle the dict->list conversion
+            input_signature_list = list(self.input_signature.values())
+
+            # Use adapted model and list signature for conversion
+            model_to_convert = adapted_model
+            signature_for_conversion = input_signature_list
+        else:
+            # No dict inputs - use model as-is
+            model_to_convert = self.model
+            signature_for_conversion = self.input_signature
+
+        # Store original model reference for later use
+        original_model = self.model
+
+        # Temporarily replace self.model with the model to convert
+        self.model = model_to_convert
+
+        try:
+            # 4. Convert the model to TFLite.
+            tflite_model = self._convert_to_tflite(signature_for_conversion)
+        finally:
+            # Restore original model
+            self.model = original_model
+
+        # 4. Save the initial TFLite model to the specified file path.
+        if not filepath.endswith(".tflite"):
+            raise ValueError(
+                "The LiteRT export requires the filepath to end with "
+                "'.tflite'. Got: {filepath}"
+            )
+
+        with open(filepath, "wb") as f:
+            f.write(tflite_model)
+
+        return filepath
+
+    def _create_dict_adapter(self, input_signature_dict):
+        """Create an adapter model that converts list inputs to dict inputs.
+
+        This adapter allows models expecting dictionary inputs to be exported
+        to TFLite format (which only supports positional/list inputs).
+
+        Args:
+            input_signature_dict: Dictionary mapping input names to InputSpec
+
+        Returns:
+            A Functional model that accepts list inputs and converts to dict
+        """
+        io_utils.print_msg(
+            f"Creating adapter for dictionary inputs: "
+            f"{list(input_signature_dict.keys())}"
+        )
+
+        input_keys = list(input_signature_dict.keys())
+
+        # Create Input layers for TFLite (list-based)
+        input_layers = []
+        for name in input_keys:
+            spec = input_signature_dict[name]
+            input_layer = layers.Input(
+                shape=spec.shape[1:],  # Remove batch dimension
+                dtype=spec.dtype,
+                name=name,
+            )
+            input_layers.append(input_layer)
+
+        # Create dict from list inputs
+        inputs_dict = {
+            name: layer for name, layer in zip(input_keys, input_layers)
+        }
+
+        # Call the original model with dict inputs
+        outputs = self.model(inputs_dict)
+
+        # Build as Functional model (list inputs -> dict -> model -> output)
+        adapted_model = models.Model(inputs=input_layers, outputs=outputs)
+
+        # Preserve the original model's variables
+        adapted_model._variables = self.model.variables
+        adapted_model._trainable_variables = self.model.trainable_variables
+        adapted_model._non_trainable_variables = (
+            self.model.non_trainable_variables
+        )
+
+        return adapted_model
+
+    def _convert_to_tflite(self, input_signature):
+        """Converts the Keras model to TFLite format.
+
+        Returns:
+            A bytes object containing the serialized TFLite model.
+        """
+        # Try direct conversion first for all models
+        try:
+            converter = tf.lite.TFLiteConverter.from_keras_model(self.model)
+            converter.target_spec.supported_ops = [
+                tf.lite.OpsSet.TFLITE_BUILTINS,
+                tf.lite.OpsSet.SELECT_TF_OPS,
+            ]
+            # Keras 3 only supports resource variables
+            converter.experimental_enable_resource_variables = True
+
+            # Apply any additional converter settings from kwargs
+            self._apply_converter_kwargs(converter)
+
+            tflite_model = converter.convert()
+
+            return tflite_model
+
+        except Exception as e:
+            # If direct conversion fails, raise the error with helpful message
+            raise RuntimeError(
+                f"Direct TFLite conversion failed. This may be due to model "
+                f"complexity or unsupported operations. Error: {e}"
+            ) from e
+
+    def _apply_converter_kwargs(self, converter):
+        """Apply additional converter settings from kwargs.
+
+        Args:
+            converter: tf.lite.TFLiteConverter instance to configure
+
+        Raises:
+            ValueError: If any kwarg is not a valid converter attribute
+        """
+        for attr, value in self.kwargs.items():
+            if attr == "target_spec" and isinstance(value, dict):
+                # Handle nested target_spec settings
+                for spec_key, spec_value in value.items():
+                    if hasattr(converter.target_spec, spec_key):
+                        setattr(converter.target_spec, spec_key, spec_value)
+                    else:
+                        raise ValueError(
+                            f"Unknown target_spec attribute '{spec_key}'"
+                        )
+            elif hasattr(converter, attr):
+                setattr(converter, attr, value)
+            else:
+                raise ValueError(f"Unknown converter attribute '{attr}'")