From a67bcd6eccaa4de00a08a3def8323debecff84d7 Mon Sep 17 00:00:00 2001
From: Madeleine Dunn <madeleine.dunn@arm.com>
Date: Thu, 5 Oct 2023 18:19:23 +0100
Subject: feat: Implement sparsity

- Add a placeholder file and registry option for sparsity

Resolves: MLIA-999

Signed-off-by: Madeleine Dunn <madeleine.dunn@arm.com>
Change-Id: I273192ba6813309f5226e3d8e0b686ce87ee6b79
---
 src/mlia/nn/rewrite/core/rewrite.py                |  6 +++++-
 src/mlia/nn/rewrite/library/fc_sparsity24_layer.py | 16 ++++++++++++++++
 tests/test_cli_commands.py                         |  3 ++-
 tests/test_nn_rewrite_core_rewrite.py              |  5 ++++-
 4 files changed, 27 insertions(+), 3 deletions(-)
 create mode 100644 src/mlia/nn/rewrite/library/fc_sparsity24_layer.py

diff --git a/src/mlia/nn/rewrite/core/rewrite.py b/src/mlia/nn/rewrite/core/rewrite.py
index f5f5561..03af0e9 100644
--- a/src/mlia/nn/rewrite/core/rewrite.py
+++ b/src/mlia/nn/rewrite/core/rewrite.py
@@ -115,7 +115,11 @@ class RewritingOptimizer(Optimizer):
         [
             DynamicallyLoadedRewrite(
                 "fully-connected", "mlia.nn.rewrite.library.fc_layer.get_keras_model"
-            )
+            ),
+            DynamicallyLoadedRewrite(
+                "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
new file mode 100644
index 0000000..fda619f
--- /dev/null
+++ b/src/mlia/nn/rewrite/library/fc_sparsity24_layer.py
@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright 2023, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""PLACEHOLDER for example rewrite with one fully connected 2:4 sparsity layer."""
+from typing import Any
+
+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."""
+    model = get_keras_model(input_shape, output_shape)
+    return model
diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py
index 9cda27c..5b2b947 100644
--- a/tests/test_cli_commands.py
+++ b/tests/test_cli_commands.py
@@ -126,7 +126,8 @@ def test_performance_unknown_target(
                 Exception,
                 match=re.escape(
                     "Invalid rewrite target: 'random'. "
-                    "Supported rewrites: ['fully-connected']"
+                    "Supported rewrites: ['fully-connected',"
+                    " 'fully_connected_sparsity24']"
                 ),
             ),
         ],
diff --git a/tests/test_nn_rewrite_core_rewrite.py b/tests/test_nn_rewrite_core_rewrite.py
index b32fafd..d9c74ea 100644
--- a/tests/test_nn_rewrite_core_rewrite.py
+++ b/tests/test_nn_rewrite_core_rewrite.py
@@ -100,7 +100,10 @@ def test_register_rewrite_function() -> None:
 
 def test_builtin_rewrite_names() -> None:
     """Test if all builtin rewrites are properly registered and returned."""
-    assert RewritingOptimizer.builtin_rewrite_names() == ["fully-connected"]
+    assert RewritingOptimizer.builtin_rewrite_names() == [
+        "fully-connected",
+        "fully_connected_sparsity24",
+    ]
 
 
 def test_rewrite_function_autoload() -> None:
-- 
cgit v1.2.1