From 2973b6d52914023f9b8797aec8309957457d4189 Mon Sep 17 00:00:00 2001
From: Nathan Bailey <nathan.bailey@arm.com>
Date: Mon, 19 Feb 2024 14:54:51 +0000
Subject: feat: CLI and API changes for the clustering rewrite

Adds API changes for a fully-connected-clustering rewrite

Resolves: MLIA-1077

Signed-off-by: Nathan Bailey <nathan.bailey@arm.com>
Change-Id: I845796a391c5020e66472456b97ecad5ee8139a8
---
 src/mlia/nn/rewrite/core/rewrite.py                |  1 +
 src/mlia/nn/rewrite/library/fc_clustering_layer.py | 13 +++++++++++++
 tests/test_cli_commands.py                         | 15 ++++++++++++++-
 tests/test_nn_rewrite_core_rewrite.py              |  2 ++
 4 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 src/mlia/nn/rewrite/library/fc_clustering_layer.py

diff --git a/src/mlia/nn/rewrite/core/rewrite.py b/src/mlia/nn/rewrite/core/rewrite.py
index 4fe1c26..a8084e8 100644
--- a/src/mlia/nn/rewrite/core/rewrite.py
+++ b/src/mlia/nn/rewrite/core/rewrite.py
@@ -157,6 +157,7 @@ class RewritingOptimizer(Optimizer):
         [
             FullyConnectedRewrite("fully-connected", fc_rewrite),
             Sparsity24Rewrite("fully-connected-sparsity24", fc_rewrite_sparsity24),
+            FullyConnectedRewrite("fully-connected-clustering", fc_rewrite),
         ]
     )
 
diff --git a/src/mlia/nn/rewrite/library/fc_clustering_layer.py b/src/mlia/nn/rewrite/library/fc_clustering_layer.py
new file mode 100644
index 0000000..07c07ac
--- /dev/null
+++ b/src/mlia/nn/rewrite/library/fc_clustering_layer.py
@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Example rewrite with one fully connected clustered layer."""
+from typing import Any
+
+from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
+
+from mlia.nn.rewrite.library.fc_layer import get_keras_model
+
+
+def get_keras_model_clus(input_shape: Any, output_shape: Any) -> keras.Model:
+    """Generate TensorFlow Lite model for clustering rewrite."""
+    return get_keras_model(input_shape, output_shape)
diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py
index e4b2609..93a05bd 100644
--- a/tests/test_cli_commands.py
+++ b/tests/test_cli_commands.py
@@ -140,7 +140,7 @@ def test_performance_unknown_target(
                 match=re.escape(
                     "Invalid rewrite target: 'random'. "
                     "Supported rewrites: ['fully-connected',"
-                    " 'fully-connected-sparsity24']"
+                    " 'fully-connected-clustering', 'fully-connected-sparsity24']"
                 ),
             ),
         ],
@@ -182,6 +182,19 @@ def test_performance_unknown_target(
                 ),
             ),
         ],
+        [
+            "ethos-u55-256",
+            False,
+            False,
+            None,
+            None,
+            None,
+            True,
+            "fully-connected-clustering",
+            "sequential/flatten/Reshape",
+            "StatefulPartitionedCall:0",
+            does_not_raise(),
+        ],
     ],
 )
 def test_opt_valid_optimization_target(  # pylint: disable=too-many-locals,too-many-arguments
diff --git a/tests/test_nn_rewrite_core_rewrite.py b/tests/test_nn_rewrite_core_rewrite.py
index 8ef5bd2..96e4160 100644
--- a/tests/test_nn_rewrite_core_rewrite.py
+++ b/tests/test_nn_rewrite_core_rewrite.py
@@ -63,6 +63,7 @@ def test_rewrite_selection(
     [
         ("fully-connected", does_not_raise()),
         ("fully-connected-sparsity24", does_not_raise()),
+        ("fully-connected-clustering", does_not_raise()),
         ("random", does_not_raise()),
     ],
 )
@@ -124,6 +125,7 @@ def test_builtin_rewrite_names() -> None:
     """Test if all builtin rewrites are properly registered and returned."""
     assert RewritingOptimizer.builtin_rewrite_names() == [
         "fully-connected",
+        "fully-connected-clustering",
         "fully-connected-sparsity24",
     ]
 
-- 
cgit v1.2.1