Skip to content

Commit

Permalink
[Stable APIs] Create torchlib_opset for torch 2.6 (#1963)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
justinchuby authored Nov 23, 2024
1 parent 99b3d26 commit 9592227
Showing 1 changed file with 19 additions and 0 deletions.
19 changes: 19 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,24 @@
save_model_with_external_data,
)

if TYPE_CHECKING:
from onnxscript.onnx_opset._impl.opset18 import Opset18


def optimize(model: ir.Model) -> ir.Model:
"""Optimize the model."""
optimizer.optimize_ir(model)
return model


def torchlib_opset() -> Opset18:
"""Return the default opset for torchlib."""
import onnxscript # pylint: disable=import-outside-toplevel

return onnxscript.opset18 # type: ignore


def torchlib_opset_version() -> int:
"""Return the default opset version for torchlib."""

return torchlib_opset().version

0 comments on commit 9592227

Please sign in to comment.