From 5216a1884352a21dabfadac4c9ec1425fe256c4b Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 22 Nov 2024 14:44:30 -0800
Subject: [PATCH] [Stable APIs] Create torchlib_opset for torch 2.6

---
 onnxscript/_framework_apis/torch_2_6.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

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