fix: Change training_parameters to return empty list instead of list of None if needed.

Extension to MLIA-1004

Signed-off-by: Nathan Bailey <nathan.bailey@arm.com>
Change-Id: Ib40c2e5932c1210a1d141200815a76e33f5ab078

2 files changed, 15 insertions, 21 deletions

diff --git a/src/mlia/nn/select.py b/src/mlia/nn/select.py
index 81a614f..b61e713 100644
--- a/src/mlia/nn/select.py
+++ b/src/mlia/nn/select.py
@@ -117,7 +117,7 @@ class MultiStageOptimizer(Optimizer):
 def get_optimizer(
     model: keras.Model | KerasModel | TFLiteModel,
     config: OptimizerConfiguration | OptimizationSettings | list[OptimizationSettings],
-    training_parameters: list[dict | None] | None = None,
+    training_parameters: dict | None = None,
 ) -> Optimizer:
     """Get optimizer for provided configuration."""
     if isinstance(model, KerasModel):
@@ -151,7 +151,7 @@ def get_optimizer(
 def _get_optimizer(
     model: keras.Model | Path,
     optimization_settings: OptimizationSettings | list[OptimizationSettings],
-    training_parameters: list[dict | None] | None = None,
+    training_parameters: dict | None = None,
 ) -> Optimizer:
     if isinstance(optimization_settings, OptimizationSettings):
         optimization_settings = [optimization_settings]
@@ -173,22 +173,17 @@ def _get_optimizer(
 
 
 def _get_rewrite_params(
-    training_parameters: list[dict | None] | None = None,
-) -> list:
+    training_parameters: dict | None = None,
+) -> TrainingParameters:
     """Get the rewrite TrainingParameters.
 
     Return the default constructed TrainingParameters() per default, but can be
     overwritten in the unit tests.
     """
-    if training_parameters is None:
-        return [TrainingParameters()]
+    if not training_parameters:
+        return TrainingParameters()
 
-    if training_parameters[0] is None:
-        train_params = TrainingParameters()
-    else:
-        train_params = TrainingParameters(**training_parameters[0])
-
-    return [train_params]
+    return TrainingParameters(**training_parameters)
 
 
 def _get_optimizer_configuration(
@@ -196,7 +191,7 @@ def _get_optimizer_configuration(
     optimization_target: int | float | str,
     layers_to_optimize: list[str] | None = None,
     dataset: Path | None = None,
-    training_parameters: list[dict | None] | None = None,
+    training_parameters: dict | None = None,
 ) -> OptimizerConfiguration:
     """Get optimizer configuration for provided parameters."""
     _check_optimizer_params(optimization_type, optimization_target)
@@ -222,7 +217,7 @@ def _get_optimizer_configuration(
                 optimization_target=str(optimization_target),
                 layers_to_optimize=layers_to_optimize,
                 dataset=dataset,
-                train_params=rewrite_params[0],
+                train_params=rewrite_params,
             )
 
         raise ConfigurationError(
diff --git a/src/mlia/target/common/optimization.py b/src/mlia/target/common/optimization.py
index 8c5d184..1423189 100644
--- a/src/mlia/target/common/optimization.py
+++ b/src/mlia/target/common/optimization.py
@@ -86,7 +86,7 @@ class OptimizingDataCollector(ContextAwareDataCollector):
     def optimize_model(
         self,
         opt_settings: list[OptimizationSettings],
-        training_parameters: list[dict | None],
+        training_parameters: dict | None,
         model: KerasModel | TFLiteModel,
     ) -> Any:
         """Run optimization."""
@@ -123,12 +123,12 @@ class OptimizingDataCollector(ContextAwareDataCollector):
             context=context,
         )
 
-    def _get_training_settings(self, context: Context) -> list[dict]:
+    def _get_training_settings(self, context: Context) -> dict:
         """Get optimization settings."""
         return self.get_parameter(  # type: ignore
             OptimizingDataCollector.name(),
             "training_parameters",
-            expected_type=list,
+            expected_type=dict,
             expected=False,
             context=context,
         )
@@ -228,9 +228,8 @@ def add_common_optimization_params(advisor_parameters: dict, extra_args: dict) -
         raise TypeError("Optimization targets value has wrong format.")
 
     rewrite_parameters = extra_args.get("optimization_profile")
-    if not rewrite_parameters:
-        training_parameters = None
-    else:
+    training_parameters = None
+    if rewrite_parameters:
         if not isinstance(rewrite_parameters, dict):
             raise TypeError("Training Parameter values has wrong format.")
         training_parameters = extra_args["optimization_profile"].get("training")
@@ -239,7 +238,7 @@ def add_common_optimization_params(advisor_parameters: dict, extra_args: dict) -
         {
             "common_optimizations": {
                 "optimizations": [optimization_targets],
-                "training_parameters": [training_parameters],
+                "training_parameters": training_parameters,
             },
         }
     )