Add Op (Slice - complex) | feat torchlib

onnxscript/function_libs/torch_lib/ops/core.py

@@ -7698,6 +7698,21 @@
     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,

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,