aboutsummaryrefslogtreecommitdiff
path: root/tests
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
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')
-rw-r--r--tests/test_cli_commands.py15
-rw-r--r--tests/test_nn_rewrite_core_rewrite.py2
2 files changed, 16 insertions, 1 deletions
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",
]