5 files changed, 46 insertions, 18 deletions

diff --git a/src/mlia/target/ethos_u/advice_generation.py b/src/mlia/target/ethos_u/advice_generation.py
index a9f9eac..351082a 100644
--- a/src/mlia/target/ethos_u/advice_generation.py
+++ b/src/mlia/target/ethos_u/advice_generation.py
@@ -11,7 +11,7 @@ from mlia.core.advice_generation import ContextAwareAdviceProducer
 from mlia.core.advice_generation import FactBasedAdviceProducer
 from mlia.core.common import AdviceCategory
 from mlia.core.common import DataItem
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.select import OptimizationSettings
 from mlia.target.common.reporters import handle_model_is_not_tflite_compatible_common
 from mlia.target.common.reporters import handle_tflite_check_failed_common
 from mlia.target.common.reporters import ModelIsNotTFLiteCompatible
diff --git a/src/mlia/target/ethos_u/advisor.py b/src/mlia/target/ethos_u/advisor.py
index d2c308a..321734c 100644
--- a/src/mlia/target/ethos_u/advisor.py
+++ b/src/mlia/target/ethos_u/advisor.py
@@ -54,8 +54,20 @@ class EthosUInferenceAdvisor(DefaultInferenceAdvisor):
         if is_tflite_model(model):
             # TensorFlow Lite models do not support optimization (only performance)!
             if context.category_enabled(AdviceCategory.OPTIMIZATION):
-                raise RuntimeError(
-                    "Optimizations are not supported for TensorFlow Lite files."
+                optimization_settings = self._get_optimization_settings(context)
+
+                optimization_types = {
+                    opt["optimization_type"] for opt in optimization_settings[0]
+                }
+                if optimization_types != {"rewrite"}:
+                    raise RuntimeError(
+                        "Only 'rewrite' is supported for TensorFlow Lite files."
+                    )
+
+                collectors.append(
+                    EthosUOptimizationPerformance(
+                        model, target_config, optimization_settings, backends
+                    )
                 )
             if context.category_enabled(AdviceCategory.PERFORMANCE):
                 collectors.append(EthosUPerformance(model, target_config, backends))
diff --git a/src/mlia/target/ethos_u/data_analysis.py b/src/mlia/target/ethos_u/data_analysis.py
index 3df4bff..5c6080f 100644
--- a/src/mlia/target/ethos_u/data_analysis.py
+++ b/src/mlia/target/ethos_u/data_analysis.py
@@ -10,7 +10,7 @@ from mlia.backend.vela.compat import Operators
 from mlia.core.common import DataItem
 from mlia.core.data_analysis import Fact
 from mlia.core.data_analysis import FactExtractor
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.select import OptimizationSettings
 from mlia.nn.tensorflow.tflite_compat import TFLiteCompatibilityInfo
 from mlia.target.common.reporters import analyze_tflite_compatibility_common
 from mlia.target.ethos_u.performance import OptimizationPerformanceMetrics
diff --git a/src/mlia/target/ethos_u/data_collection.py b/src/mlia/target/ethos_u/data_collection.py
index 0654143..0f3a8d2 100644
--- a/src/mlia/target/ethos_u/data_collection.py
+++ b/src/mlia/target/ethos_u/data_collection.py
@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import logging
 from pathlib import Path
+from typing import Any
 
 from mlia.backend.vela.compat import Operators
 from mlia.backend.vela.compat import supported_operators
@@ -12,15 +13,17 @@ from mlia.core.context import Context
 from mlia.core.data_collection import ContextAwareDataCollector
 from mlia.core.errors import FunctionalityNotSupportedError
 from mlia.core.performance import estimate_performance
+from mlia.nn.select import get_optimizer
+from mlia.nn.select import OptimizationSettings
 from mlia.nn.tensorflow.config import get_keras_model
 from mlia.nn.tensorflow.config import get_tflite_model
 from mlia.nn.tensorflow.config import KerasModel
-from mlia.nn.tensorflow.optimizations.select import get_optimizer
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
+from mlia.nn.tensorflow.config import TFLiteModel
 from mlia.nn.tensorflow.tflite_compat import TFLiteChecker
 from mlia.nn.tensorflow.tflite_compat import TFLiteCompatibilityInfo
 from mlia.nn.tensorflow.utils import is_tflite_model
 from mlia.nn.tensorflow.utils import save_keras_model
+from mlia.nn.tensorflow.utils import save_tflite_model
 from mlia.target.ethos_u.config import EthosUConfiguration
 from mlia.target.ethos_u.performance import EthosUPerformanceEstimator
 from mlia.target.ethos_u.performance import OptimizationPerformanceMetrics
@@ -103,7 +106,7 @@ class OptimizeModel:
         self.context = context
         self.opt_settings = opt_settings
 
-    def __call__(self, keras_model: KerasModel) -> KerasModel:
+    def __call__(self, keras_model: KerasModel) -> Any:
         """Run optimization."""
         optimizer = get_optimizer(keras_model, self.opt_settings)
 
@@ -112,9 +115,19 @@ class OptimizeModel:
         optimizer.apply_optimization()
 
         model = optimizer.get_model()
+
+        if isinstance(model, Path):
+            return model
+
+        if isinstance(model, TFLiteModel):
+            model_path = self.context.get_model_path("optimized_model.tflite")
+            with open(model.model_path, "rb") as file_handle:
+                model_data = bytearray(file_handle.read())
+            save_tflite_model(model_data, model_path)
+            return TFLiteModel(model_path)
+
         model_path = self.context.get_model_path("optimized_model.h5")
         save_keras_model(model, model_path)
-
         return KerasModel(model_path)
 
 
@@ -146,14 +159,17 @@ class EthosUOptimizationPerformance(ContextAwareDataCollector):
 
         opt_settings = self._parse_optimization_params(self.optimizations)
 
-        try:
-            keras_model = get_keras_model(self.model, self.context)
-        except NotImplementedError as err:
-            raise FunctionalityNotSupportedError(
-                reason="Unable to run model optimizations",
-                description=f"{self.model} is not a Keras model and "
-                "could not be converted to a Keras model",
-            ) from err
+        if opt_settings[0][0].optimization_type != "rewrite":
+            try:
+                model = get_keras_model(self.model, self.context)
+            except NotImplementedError as err:
+                raise FunctionalityNotSupportedError(
+                    reason="Unable to run model optimizations",
+                    description=f"{self.model} is not a Keras model and "
+                    "could not be converted to a Keras model",
+                ) from err
+        else:
+            model = self.model  # type: ignore
 
         optimizers = [OptimizeModel(self.context, opts) for opts in opt_settings]
 
@@ -163,7 +179,7 @@ class EthosUOptimizationPerformance(ContextAwareDataCollector):
             self.backends,
         )
         original_metrics, *optimized_metrics = estimate_performance(
-            keras_model, estimator, optimizers  # type: ignore
+            model, estimator, optimizers  # type: ignore
         )
 
         result = OptimizationPerformanceMetrics(
diff --git a/src/mlia/target/ethos_u/performance.py b/src/mlia/target/ethos_u/performance.py
index f7f9a8c..a0526e4 100644
--- a/src/mlia/target/ethos_u/performance.py
+++ b/src/mlia/target/ethos_u/performance.py
@@ -15,9 +15,9 @@ from mlia.backend.corstone import is_corstone_backend
 from mlia.backend.corstone.performance import estimate_performance
 from mlia.core.context import Context
 from mlia.core.performance import PerformanceEstimator
+from mlia.nn.select import OptimizationSettings
 from mlia.nn.tensorflow.config import get_tflite_model
 from mlia.nn.tensorflow.config import ModelConfiguration
-from mlia.nn.tensorflow.optimizations.select import OptimizationSettings
 from mlia.target.ethos_u.config import EthosUConfiguration
 from mlia.target.registry import supported_backends
 from mlia.utils.logging import log_action