Skip to content

Commit 9592227

Browse files
authored
[Stable APIs] Create torchlib_opset for torch 2.6 (microsoft#1963)
Create torchlib_opset for torch 2.6. This will be used for creating the model opset import as well as in `_building` for creating constant/concat nodes etc.
1 parent 99b3d26 commit 9592227

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

onnxscript/_framework_apis/torch_2_6.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,10 @@
1010
"get_torchlib_ops",
1111
"optimize",
1212
"save_model_with_external_data",
13+
"torchlib_opset",
1314
]
15+
from typing import TYPE_CHECKING
16+
1417
from onnxscript import ir, optimizer
1518
from onnxscript._framework_apis.torch_2_5 import (
1619
check_model,
@@ -19,8 +22,24 @@
1922
save_model_with_external_data,
2023
)
2124

25+
if TYPE_CHECKING:
26+
from onnxscript.onnx_opset._impl.opset18 import Opset18
27+
2228

2329
def optimize(model: ir.Model) -> ir.Model:
2430
"""Optimize the model."""
2531
optimizer.optimize_ir(model)
2632
return model
33+
34+
35+
def torchlib_opset() -> Opset18:
36+
"""Return the default opset for torchlib."""
37+
import onnxscript # pylint: disable=import-outside-toplevel
38+
39+
return onnxscript.opset18 # type: ignore
40+
41+
42+
def torchlib_opset_version() -> int:
43+
"""Return the default opset version for torchlib."""
44+
45+
return torchlib_opset().version

0 commit comments

Comments
 (0)