unit tests

neuralmagic · Jun 13, 2023 · 92b25ff · 92b25ff
1 parent 177a4b3
commit 92b25ff
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/tests/sparseml/exporters/transforms/test_matmul_add_to_matmulinteger_add_cast_mul.py b/tests/sparseml/exporters/transforms/test_matmul_add_to_matmulinteger_add_cast_mul.py
@@ -146,25 +146,22 @@ def test_without_transpose(onnx_model: onnx.ModelProto):
     ]
 
 
-def test_no_bias_changes_nothing(onnx_model: onnx.ModelProto):
+def test_matmul_no_bias_converts(onnx_model: onnx.ModelProto):
     # remove "bias" initializer and "add" node
     assert onnx_model.graph.initializer.pop().name == "bias"
     assert onnx_model.graph.node.pop().name == "add"
     validate_onnx(onnx_model)
 
     onnx_model = MatMulAddToMatMulIntegerAddCastMul().apply(onnx_model)
     validate_onnx(onnx_model)
-    # NOTE: nothing changes
+    # converted model should have matmulinteger + rescale mul without bias add
     assert [i.name for i in onnx_model.graph.initializer] == [
-        "x_scale",
-        "y_scale",
         "zero_point",
-        "weight",
+        "matmul.weight_quantized",
+        "matmul_quant.rescale.scale",
     ]
     assert [n.name for n in onnx_model.graph.node] == [
-        "input_dequant",
-        "weight_quant",
-        "weight_dequant",
-        "transpose",
-        "matmul",
+        "matmul_quant",
+        "matmul_bias_add_quant_cast",
+        "matmul_quant_rescale_mul",
     ]