File tree 1 file changed +9
-0
lines changed
onnxscript/function_libs/torch_lib
1 file changed +9
-0
lines changed Original file line number Diff line number Diff line change 17
17
from typing_extensions import TypeAlias
18
18
19
19
import onnxscript
20
+
21
+ # Q: Don't know why vscode only recognizes this import style instead of `from onnxscript import optimizer`.
22
+ import onnxscript .optimizer as optimizer
20
23
from onnxscript import evaluator
21
24
from onnxscript import tensor as onnxscript_tensor
22
25
from onnxscript ._internal import param_manipulation , runtime_typing
23
26
from onnxscript .function_libs .torch_lib import _flags
24
27
from onnxscript .function_libs .torch_lib .ops import common as common_ops
28
+ from onnxscript .rewriter import onnxruntime as ort_rewriter
25
29
26
30
__all__ = [
27
31
"TorchScriptTensor" ,
@@ -1071,4 +1075,9 @@ def to_model_proto(
1071
1075
common_ops .common_opset .domain , common_ops .common_opset .version
1072
1076
)
1073
1077
)
1078
+
1079
+ # Not the best integration point. Enables benchmarking the migration.
1080
+ onnx_model = optimizer .optimize (onnx_model )
1081
+ # This also creates contrib op in the model. So definitely not the best integration point.
1082
+ onnx_model = ort_rewriter .rewrite (onnx_model )
1074
1083
return onnx_model
You can’t perform that action at this time.
0 commit comments