diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/test_cli_commands.py | 3 | ||||
| -rw-r--r-- | tests/test_nn_rewrite_core_rewrite.py | 5 |
2 files changed, 6 insertions, 2 deletions
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: |