From 97e2cdbfe52f1ad00024aa472ab4eac4d8a5f6d9 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 2 Sep 2025 10:57:11 -0700 Subject: [PATCH 1/2] [torchlib] Modify aten_unbind to use None for split_sizes --- onnxscript/function_libs/torch_lib/ops/core.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index ab992e0580..32ac1e32f6 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8616,12 +8616,11 @@ def aten_type_as(self: TTensor, other: TTensor2) -> TTensor2: return op.CastLike(self, other) -@torch_op("aten::unbind.int") +@torch_op("aten::unbind.int", trace_only=True) def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]: """unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]""" - split_sizes = op.Constant(value_int=1) - return op.SplitToSequence(self, split_sizes, axis=dim, keepdims=False) + return op.SplitToSequence(self, None, axis=dim, keepdims=False) @torch_op("aten::unflatten.int", trace_only=True) From fbfee996e781b73f782ca39560a22cf632ff925f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 2 Sep 2025 10:59:35 -0700 Subject: [PATCH 2/2] Update core.py --- onnxscript/function_libs/torch_lib/ops/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 32ac1e32f6..88c353e777 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8620,7 +8620,7 @@ 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)[]""" - return op.SplitToSequence(self, None, axis=dim, keepdims=False) + return op.SplitToSequence(self, axis=dim, keepdims=False) @torch_op("aten::unflatten.int", trace_only=True)