fix(torchlib): aten::unbind.int uses SplitToSequence(keepdims=False) to match PyTorch shapes

Author: MRADULTRIPATHI

docs/test/test_documentation_examples.py

@@ -56,3 +56,26 @@ def test(*relpath):
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
+
+#changed code
+import torch
+from onnxscript.function_libs.torch_lib import ops
+from onnxscript import evaluator
+
+def test_unbind_matches_torch():
+    x_torch = torch.randn(3, 4)
+    y_torch = torch.unbind(x_torch, dim=1)
+
+    # Convert input to NumPy for ONNXScript
+    x_np = x_torch.detach().cpu().numpy()
+
+    # Run in eager mode
+    eager = evaluator.default()
+    y_onnx = eager.eval_function(ops.core.aten_unbind, (x_np,), {"dim": 1})
+
+    # Compare number of outputs
+    assert len(y_torch) == len(y_onnx)
+
+    # Compare shapes
+    for a, b in zip(y_torch, y_onnx):
+        assert a.shape == tuple(b.shape), f"Shape mismatch: {a.shape} vs {b.shape}"

onnxscript/function_libs/torch_lib/ops/core.py

@@ -8618,10 +8618,13 @@ def aten_type_as(self: TTensor, other: TTensor2) -> TTensor2:
 
 @torch_op("aten::unbind.int")
 def aten_unbind(self: TTensor, dim: int = 0) -> Sequence[TTensor]:
-    """unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[]"""
+    # Use opset18 so keepdims=False works properly
+    shape = op.Shape(self)
+    dim_size = op.Gather(shape, op.Constant(value_int=dim), axis=0)
 
-    split_sizes = op.Constant(value_int=1)
-    return op.SplitToSequence(self, split_sizes, axis=dim, keepdims=False)
+    # Split into slices of size 1 along `dim`, dropping the axis
+    split_size = op.Constant(value_int=1)
+    return op.SplitToSequence(self, split_size, axis=dim, keepdims=False)
 
 
 @torch_op("aten::unflatten.int", trace_only=True)