From 2a2a910d6d7cc3e7555b0a3c1ba458a4065c41ae Mon Sep 17 00:00:00 2001
From: Madeleine Dunn <madeleine.dunn@arm.com>
Date: Mon, 9 Oct 2023 17:02:54 +0100
Subject: feat: Add placeholder rewrite_target for sparsity24

- The placeholder currently duplicates the existing fc target

Resolves: MLIA-1000

Signed-off-by: Madeleine Dunn <madeleine.dunn@arm.com>
Change-Id: I0df5d47e61dafa567e212566bbcb0f1639fe7642
---
 src/mlia/nn/rewrite/core/rewrite.py                |  5 +++--
 src/mlia/nn/rewrite/library/fc_sparsity24_layer.py |  2 --
 tests/test_cli_commands.py                         | 15 ++++++++++++++-
 tests/test_nn_rewrite_core_rewrite.py              |  6 ++++--
 4 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/mlia/nn/rewrite/core/rewrite.py b/src/mlia/nn/rewrite/core/rewrite.py
index 03af0e9..a4d47c4 100644
--- a/src/mlia/nn/rewrite/core/rewrite.py
+++ b/src/mlia/nn/rewrite/core/rewrite.py
@@ -114,10 +114,11 @@ class RewritingOptimizer(Optimizer):
     registry = RewriteRegistry(
         [
             DynamicallyLoadedRewrite(
-                "fully-connected", "mlia.nn.rewrite.library.fc_layer.get_keras_model"
+                "fully-connected",
+                "mlia.nn.rewrite.library.fc_layer.get_keras_model",
             ),
             DynamicallyLoadedRewrite(
-                "fully_connected_sparsity24",
+                "fully-connected-sparsity24",
                 "mlia.nn.rewrite.library.fc_sparsity24_layer.get_keras_model24",
             ),
         ]
diff --git a/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py b/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py
index fda619f..1b17522 100644
--- a/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py
+++ b/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py
@@ -7,8 +7,6 @@ from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
 from .fc_layer import get_keras_model
 
-# get_keras_model(input_shape,output_shape)
-
 
 def get_keras_model24(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for rewrite."""
diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py
index 5b2b947..e4b2609 100644
--- a/tests/test_cli_commands.py
+++ b/tests/test_cli_commands.py
@@ -82,6 +82,19 @@ def test_performance_unknown_target(
             "StatefulPartitionedCall:0",
             does_not_raise(),
         ],
+        [
+            "ethos-u55-256",
+            False,
+            False,
+            None,
+            None,
+            None,
+            True,
+            "fully-connected-sparsity24",
+            "sequential/flatten/Reshape",
+            "StatefulPartitionedCall:0",
+            does_not_raise(),
+        ],
         [
             "ethos-u55-256",
             True,
@@ -127,7 +140,7 @@ def test_performance_unknown_target(
                 match=re.escape(
                     "Invalid rewrite target: 'random'. "
                     "Supported rewrites: ['fully-connected',"
-                    " 'fully_connected_sparsity24']"
+                    " 'fully-connected-sparsity24']"
                 ),
             ),
         ],
diff --git a/tests/test_nn_rewrite_core_rewrite.py b/tests/test_nn_rewrite_core_rewrite.py
index d9c74ea..1d0100a 100644
--- a/tests/test_nn_rewrite_core_rewrite.py
+++ b/tests/test_nn_rewrite_core_rewrite.py
@@ -42,13 +42,15 @@ def test_rewrite() -> None:
     "rewrite_name, expected_error",
     [
         ("fully-connected", does_not_raise()),
+        ("fully-connected-sparsity24", does_not_raise()),
         ("random", does_not_raise()),
     ],
 )
 def test_rewrite_configuration(
     test_tflite_model_fp32: Path, rewrite_name: str, expected_error: Any
 ) -> None:
-    """Test get_rewrite function only supports rewrite type fully-connected."""
+    """Test get_rewrite function only supports rewrite types
+    fully-connected and fully-connected-sparsity24."""
     with expected_error:
         config_obj = RewriteConfiguration(
             rewrite_name,
@@ -102,7 +104,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_sparsity24",
+        "fully-connected-sparsity24",
     ]
 
 
-- 
cgit v1.2.1