Skip to content

Commit

Permalink
[Stable APIs] Create torchlib_opset for torch 2.6
Browse files Browse the repository at this point in the history
  • Loading branch information
justinchuby committed Nov 22, 2024
1 parent 5a4d22e commit 5216a18
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions onnxscript/_framework_apis/torch_2_6.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -19,8 +22,18 @@
save_model_with_external_data,
)

if TYPE_CHECKING:
from onnxscript.values import Opset

Check warning on line 26 in onnxscript/_framework_apis/torch_2_6.py

View check run for this annotation

Codecov / codecov/patch

onnxscript/_framework_apis/torch_2_6.py#L26

Added line #L26 was not covered by tests


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

Check warning on line 37 in onnxscript/_framework_apis/torch_2_6.py

View check run for this annotation

Codecov / codecov/patch

onnxscript/_framework_apis/torch_2_6.py#L37

Added line #L37 was not covered by tests

return opset18

Check warning on line 39 in onnxscript/_framework_apis/torch_2_6.py

View check run for this annotation

Codecov / codecov/patch

onnxscript/_framework_apis/torch_2_6.py#L39

Added line #L39 was not covered by tests

0 comments on commit 5216a18

Please sign in to comment.