Add Op (Slice - complex) | feat torchlib (#2089)

Fix pytorch/pytorch#147896
microsoft · Mar 7, 2025 · 68d4b9f · 68d4b9f
1 parent 4c1cda2
commit 68d4b9f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -7698,6 +7698,21 @@ def aten_sinh(self: TFloat) -> TFloat:
     return op.Sinh(self)
 
 
+@torch_op(("aten::slice.Tensor"), trace_only=True, complex=True)
+def aten_slice_complex(
+    self: TTensor,
+    dim: int = 0,
+    start: Optional[INT64] = None,
+    end: Optional[INT64] = None,
+    step: Optional[INT64] = None,
+) -> TTensor:
+    """slice.Tensor(Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a)"""
+    if dim < 0:
+        # Account for the complex dimension in ONNX
+        dim = dim - 1
+    return aten_slice(self, dim, start, end, step)
+
+
 @torch_op(("aten::slice.Tensor"), trace_only=True)
 def aten_slice(
     self: TTensor,

diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py
@@ -2048,6 +2048,7 @@ def _where_input_wrangler(
     ),
     TorchLibOpInfo("ops.aten.slice_scatter", core_ops.aten_slice_scatter),
     TorchLibOpInfo("slice", core_ops.aten_slice),
+    TorchLibOpInfo("slice", core_ops.aten_slice_complex, complex=True),
     TorchLibOpInfo(
         "sum",
         core_ops.aten_sum_dim_IntList,