-
Notifications
You must be signed in to change notification settings - Fork 56
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Stable APIs] Create torchlib_opset for torch 2.6
- Loading branch information
1 parent
5a4d22e
commit 5216a18
Showing
1 changed file
with
13 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -10,7 +10,10 @@ | |
"get_torchlib_ops", | ||
"optimize", | ||
"save_model_with_external_data", | ||
"torchlib_opset", | ||
] | ||
from typing import TYPE_CHECKING | ||
|
||
from onnxscript import ir, optimizer | ||
from onnxscript._framework_apis.torch_2_5 import ( | ||
check_model, | ||
|
@@ -19,8 +22,18 @@ | |
save_model_with_external_data, | ||
) | ||
|
||
if TYPE_CHECKING: | ||
from onnxscript.values import Opset | ||
|
||
|
||
def optimize(model: ir.Model) -> ir.Model: | ||
"""Optimize the model.""" | ||
optimizer.optimize_ir(model) | ||
return model | ||
|
||
|
||
def torchlib_opset() -> Opset: | ||
"""Return the default opset for torchlib.""" | ||
from onnxscript import opset18 | ||
Check notice Code scanning / lintrunner PYLINT/C0415 Note
Import outside toplevel (onnxscript.opset18) (import-outside-toplevel)
See import-outside-toplevel. To disable, use # pylint: disable=import-outside-toplevel |
||
|
||
return opset18 | ||