Add one test failing for the optimizer after a model optimized and in…

…lined (#1465) Co-authored-by: Justin Chu <[email protected]> Co-authored-by: Ti-Tai Wang <[email protected]>
microsoft · Apr 26, 2024 · 997beb2 · 997beb2
1 parent 7dcddea
commit 997beb2
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 1 deletion.
diff --git a/testdata/dort_models/llama_forward.onnx b/testdata/dort_models/llama_forward.onnx
diff --git a/tests/optimizer/test_models.py b/tests/optimizer/test_models.py
@@ -6,10 +6,12 @@
 
 import numpy as np
 import onnx
+import onnx.inliner
 import onnxruntime
 import parameterized
 
 from onnxscript import optimizer
+from onnxscript.rewriter import onnxruntime as ort_rewriter
 from onnxscript.utils import evaluation_utils
 
 _SKIP_TABLE = {}
@@ -64,6 +66,41 @@ def test_model_runs_and_matches_accuracy_after_optimization(self, model_name):
             for output, expected_output in zip(outputs, expected_outputs):
                 np.testing.assert_allclose(output, expected_output, rtol=1e-3, atol=1e-3)
 
+    def test_optimizer_after_inlining(self):
+        model_dir = pathlib.Path(model_folder_path) / ".." / "dort_models"
+        filename = model_dir / "llama_forward.onnx"
+
+        onnx_model = onnx.load(filename)
+        onnxruntime.InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+
+        # first time
+        onnx_model = optimizer.optimize(onnx_model)
+        onnxruntime.InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        onnx_model = ort_rewriter.rewrite(onnx_model)
+        onnxruntime.InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+
+        # inline
+        onnx_model = onnx.inliner.inline_local_functions(onnx_model)
+        onnxruntime.InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+
+        # second time
+        onnx_model = optimizer.optimize(onnx_model)
+        onnxruntime.InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        onnx_model = ort_rewriter.rewrite(onnx_model)
+        onnxruntime.InferenceSession(
+            onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main(verbosity=2)