feat: Implement the clustering rewrite for int8

Implements a clustering rewrite for fully connected layers for int8 models

Resolves: MLIA-1080

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

1 files changed, 3 insertions, 9 deletions

diff --git a/tests/test_nn_rewrite_core_train.py b/tests/test_nn_rewrite_core_train.py
index 371c79f..94c99ff 100644
--- a/tests/test_nn_rewrite_core_train.py
+++ b/tests/test_nn_rewrite_core_train.py
@@ -14,15 +14,13 @@ import pytest
 import tensorflow as tf
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
-from mlia.nn.rewrite.core.rewrite import FullyConnectedRewrite
-from mlia.nn.rewrite.core.rewrite import QATRewrite
 from mlia.nn.rewrite.core.train import augment_fn_twins
 from mlia.nn.rewrite.core.train import AUGMENTATION_PRESETS
 from mlia.nn.rewrite.core.train import LearningRateSchedule
 from mlia.nn.rewrite.core.train import mixup
 from mlia.nn.rewrite.core.train import train
 from mlia.nn.rewrite.core.train import TrainingParameters
-from mlia.nn.rewrite.library.fc_layer import get_keras_model as fc_rewrite
+from tests.test_nn_rewrite_core_rewrite import TestRewrite
 from tests.utils.rewrite import MockTrainingParameters
 
 
@@ -56,20 +54,16 @@ def check_train(
     """Test the train() function."""
     with TemporaryDirectory() as tmp_dir:
         output_file = Path(tmp_dir, "out.tflite")
-        mock_rewrite = FullyConnectedRewrite(
-            name="replace",
-            rewrite_fn=fc_rewrite,
-        )
-        is_qat = isinstance(mock_rewrite, QATRewrite)
+        mock_rewrite = TestRewrite("replace", replace_fully_connected_with_conv)
         result = train(
             source_model=str(tflite_model),
             unmodified_model=str(tflite_model) if use_unmodified_model else None,
             output_model=str(output_file),
             input_tfrec=str(tfrecord),
             rewrite=mock_rewrite,
+            is_qat=False,
             input_tensors=["sequential/flatten/Reshape"],
             output_tensors=["StatefulPartitionedCall:0"],
-            is_qat=is_qat,
             train_params=train_params,
         )