Model Export to liteRT

keras/src/export/__init__.py

@@ -3,3 +3,12 @@
 from keras.src.export.saved_model import ExportArchive
 from keras.src.export.saved_model import export_saved_model
 from keras.src.export.tfsm_layer import TFSMLayer
+
+# LiteRT export requires TensorFlow, so we import conditionally
+try:
+    from keras.src.export.litert import LitertExporter
+    from keras.src.export.litert import export_litert
+except ImportError:
+    # TensorFlow not available, LiteRT export will not be available
+    LitertExporter = None
+    export_litert = None

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,18 +25,25 @@ 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(
-                "The model provided has never called. "
-                "It must be called at least once before export."
+                "The model provided has never called. It must be called "
+                "at least once before export."
             )
     return input_signature
 
@@ -41,25 +56,36 @@ def _infer_input_signature_from_model(model):
     def _make_input_spec(structure):
         # We need to turn wrapper structures like TrackingDict or _DictWrapper
         # into plain Python structures because they don't work with jax2tf/JAX.
+        if structure is None:
+            return None
         if isinstance(structure, dict):
             return {k: _make_input_spec(v) for k, v in structure.items()}
         elif isinstance(structure, tuple):
             if all(isinstance(d, (int, type(None))) for d in structure):
-                return layers.InputSpec(
-                    shape=(None,) + structure[1:], dtype=model.input_dtype
+                # For export, force batch dimension to None for flexible
+                # batching
+                shape = (
+                    (None,) + structure[1:] if len(structure) > 0 else structure
                 )
+                return layers.InputSpec(shape=shape, dtype=model.input_dtype)
             return tuple(_make_input_spec(v) for v in structure)
         elif isinstance(structure, list):
             if all(isinstance(d, (int, type(None))) for d in structure):
-                return layers.InputSpec(
-                    shape=[None] + structure[1:], dtype=model.input_dtype
+                # For export, force batch dimension to None for flexible
+                # batching
+                shape = (
+                    (None,) + tuple(structure[1:])
+                    if len(structure) > 0
+                    else tuple(structure)
                 )
+                return layers.InputSpec(shape=shape, dtype=model.input_dtype)
             return [_make_input_spec(v) for v in structure]
         else:
             raise ValueError(
                 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()]