diff --git a/onnxscript/_framework_apis/torch_2_6.py b/onnxscript/_framework_apis/torch_2_6.py index ec929a1d8..df6a3ad88 100644 --- a/onnxscript/_framework_apis/torch_2_6.py +++ b/onnxscript/_framework_apis/torch_2_6.py @@ -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 + + return opset18