aboutsummaryrefslogtreecommitdiff
path: root/tests/test_nn_rewrite_core_rewrite.py
diff options
context:
space:
mode:
authorNathan Bailey <nathan.bailey@arm.com>2024-02-19 14:54:51 +0000
committerNathan Bailey <nathan.bailey@arm.com>2024-04-16 12:31:09 +0100
commit2973b6d52914023f9b8797aec8309957457d4189 (patch)
tree5828a9d988ec2e8a95633a28e0c744ba61d9b154 /tests/test_nn_rewrite_core_rewrite.py
parentec59b3c95106daebe2ce0e57592b2bf9e6562f54 (diff)
downloadmlia-2973b6d52914023f9b8797aec8309957457d4189.tar.gz
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
Diffstat (limited to 'tests/test_nn_rewrite_core_rewrite.py')
-rw-r--r--tests/test_nn_rewrite_core_rewrite.py2
1 files changed, 2 insertions, 0 deletions
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",
]