up

Author: metascroy

backends/apple/coreml/test/test_torch_ops.py

@@ -17,7 +17,7 @@
 
 from torchao.prototype.quantization.codebook_coreml import CodebookWeightOnlyConfig
 from torchao.quantization import IntxWeightOnlyConfig, PerAxis, PerGroup, quantize_
-
+from executorch.exir.backend.utils import format_delegated_graph
 
 def is_fbcode():
     return not hasattr(torch.version, "git_version")
@@ -173,7 +173,7 @@ def test_dequantize_codebook_linear(self):
             CodebookWeightOnlyConfig(dtype=torch.uint2, block_size=[-1, 16]),
         )
         ep = torch.export.export(model, example_inputs)
-        print("ORIGINAL MODEL", ep)
+        assert "torch.ops.quant.dequantize_codebook.default" in ep.graph_module.code
         delegated_program = executorch.exir.to_edge_transform_and_lower(
             ep,
             partitioner=[self._coreml_partitioner()],
@@ -184,8 +184,11 @@ def test_dequantize_codebook_linear(self):
                     "executorch_call_delegate",
                     "getitem",
                 ], f"Got unexpected node target after delegation: {node.target.__name__}"
+
+        print(format_delegated_graph(delegated_program.exported_program().graph_module))
+
+
         et_prog = delegated_program.to_executorch()
-        print(et_prog.exported_program())
         self._compare_outputs(et_prog, model, example_inputs)
 
     def test_dequantize_codebook_embedding(self):
@@ -196,6 +199,7 @@ def test_dequantize_codebook_embedding(self):
             lambda m, fqn: isinstance(m, torch.nn.Embedding),
         )
         ep = torch.export.export(model, example_inputs)
+        assert "torch.ops.quant.dequantize_codebook.default" in ep.graph_module.code
         delegated_program = executorch.exir.to_edge_transform_and_lower(
             ep,
             partitioner=[self._coreml_partitioner()],
@@ -207,16 +211,15 @@ def test_dequantize_codebook_embedding(self):
                     "getitem",
                 ], f"Got unexpected node target after delegation: {node.target.__name__}"
         et_prog = delegated_program.to_executorch()
-        print(et_prog.exported_program())
         self._compare_outputs(et_prog, model, example_inputs)
 
 
 if __name__ == "__main__":
     test_runner = TestTorchOps()
-    test_runner.test_dequantize_affine_b4w_embedding()
-    test_runner.test_dequantize_affine_b4w_linear()
-    test_runner.test_dequantize_affine_c4w_embedding()
-    test_runner.test_dequantize_affine_c4w_linear()
-    test_runner.test_dequantize_affine_c8w_embedding_b4w_linear()
+    # test_runner.test_dequantize_affine_b4w_embedding()
+    # test_runner.test_dequantize_affine_b4w_linear()
+    # test_runner.test_dequantize_affine_c4w_embedding()
+    # test_runner.test_dequantize_affine_c4w_linear()
+    # test_runner.test_dequantize_affine_c8w_embedding_b4w_linear()
     test_runner.test_dequantize_codebook_linear()
     test_runner.test_dequantize_codebook_embedding()