wip

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -3026,10 +3026,15 @@ def aten_imag(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
+@torch_op("aten::index.Tensor")
 def aten_index(self: TTensor, indices: Sequence[INT64]) -> TTensor:
     """index.Tensor(Tensor self, Tensor?[] indices) -> Tensor"""
 
-    return op.Gather(self, indices)
+    result = self
+    for i in range(op.SequenceLength(indices)):
+        result = op.Gather(result, op.SequenceAt(indices, i))
+
+    return result
 
 
 def aten_index_add(

onnxscript/tests/function_libs/torch_lib/extra_opinfo.py

@@ -446,6 +446,22 @@ def sample_inputs_col2im(op_info, device, dtype, requires_grad, **kwargs):
         yield opinfo_core.SampleInput(tensor, args=(output_size, kernel_size), kwargs=kwargs)
 
 
+def sample_inputs_index(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info  # Unused
+    del kwargs  # Unused
+    make_arg = functools.partial(
+        torch_testing.make_tensor, dtype=dtype, device=device, requires_grad=requires_grad
+    )
+    s = 5
+    test_args = [
+        ([common_methods_invocations.index_variable(2, s, device=device)],),
+        # ([torch.tensor()],)
+    ]
+
+    for args in test_args:
+        yield opinfo_core.SampleInput(make_arg((s, s, s)), args=args)
+
+
 def sample_inputs_stft(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     del kwargs
@@ -581,9 +597,17 @@ def sample_inputs_bernoulli_p_deterministic(op_info, device, dtype, requires_gra
         skips=(),
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "aten.index.Tensor",
+        dtypes=common_dtype.all_types_and_complex_and(
+            torch.bool, torch.float16, torch.bfloat16, torch.chalf
+        ),
+        aten_name="index",
+        op=torch.ops.aten.index.Tensor,
+        sample_inputs_func=sample_inputs_index,
+    ),
     opinfo_core.OpInfo(
         "layer_norm",
-        aliases=("layer_norm",),
         aten_name="layer_norm",
         dtypes=common_dtype.floating_and_complex_types_and(torch.int64, torch.bfloat16),
         sample_inputs_func=sample_inputs_layer_norm,

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -601,6 +601,7 @@ def _where_input_wrangler(
     TorchLibOpInfo("gt", core_ops.aten_gt),
     # TorchLibOpInfo("is_same_size", core_ops.aten_is_same_size),  # no test case in OPS_DB
     # TorchLibOpInfo("is_nonzero", core_ops.aten_is_nonzero),  # no test case in OPS_DB
+    TorchLibOpInfo("aten.index.Tensor", core_ops.aten_index),
     TorchLibOpInfo(
         "index_put_bool",
         core_ops.aten_index_put_bool,