Merge branch 'main' into justinchu/consolidate-overloads

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -8467,6 +8467,12 @@ def aten_type_as(self: TTensor, other: TTensor2) -> TTensor2:
 def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]:
     """unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]"""
 
+    if isinstance(self.shape[dim], int) and not version_utils.torch_older_than("2.7"):
+        # We can create a definitive split op if the input shape is static
+        # Only torch>=2.7 supports correctly generating the correct number of outputs for Split
+        outputs = op.Split(self, axis=dim, num_outputs=self.shape[dim])
+        return [op.Squeeze(out, [dim]) for out in outputs]
+
     return op.SplitToSequence(self, axis=dim, keepdims=False)
 
 

tests/function_libs/torch_lib/ops_test.py

@@ -39,6 +39,7 @@
 from torch.utils import _pytree as pytree
 
 import onnxscript
+from onnxscript._internal import version_utils
 from tests.function_libs.torch_lib import (
     error_reproduction,
     ops_test_common,
@@ -202,7 +203,7 @@ def run_test_output_match(
                 reference_torch_outputs, _ = pytree.tree_flatten(torch_output)
                 if (
                     op.name.startswith("split")
-                    or op.name.startswith("unbind")
+                    or (op.name.startswith("unbind") and version_utils.torch_older_than("2.7"))
                     or op.name
                     in {"atleast_1d_Sequence", "atleast_2d_Sequence", "atleast_3d_Sequence"}
                 ):