[fx][acc_ops] add acc_op mapper for torch.tensor_split (#66)

Summary: Pull Request resolved: https://github.com/pytorch/fx2trt/pull/66

Reviewed By: 842974287

Differential Revision: D36126848

fbshipit-source-id: a8324bd773105be7cbdad00681cb971aa09f1822

Author: alexbeloi

test/tracer/test_acc_tracer.py

@@ -2336,6 +2336,43 @@ def forward(self, a, b, c):
         res = traced(cond, x, y)
         self.assertTrue(torch.equal(ref, res))
 
+    @parameterized.expand(
+        [
+            ("sections divisible", 2, 0),
+            ("sections indivisible", 3, 0),
+            ("indices list", [1, 3], 0),
+            ("indices tuple", (1, 3), 0),
+            ("indices tensor", torch.tensor([1, 3]), 0),
+            ("indices tensor dim1", torch.tensor([1, 3]), 1),
+            ("indices tensor dim2", torch.tensor([1, 3]), 2),
+            ("indices tensor long dim2", torch.tensor([1, 3, 5, 7]), 2),
+        ]
+    )
+    def test_tensor_split(self, _, indices_or_sections, dim):
+        """
+        Test that the tracer works for torch.tensor_split with indices and sections
+        """
+
+        class TestModule(nn.Module):
+            def __init__(self, indices_or_sections, dim):
+                super().__init__()
+                self._indices_or_sections = indices_or_sections
+                self._dim = dim
+
+            def forward(self, a):
+                return torch.tensor_split(a, self._indices_or_sections, self._dim)
+
+        m = TestModule(indices_or_sections, dim)
+        a = torch.randn(4, 8, 16)
+        traced = acc_tracer.trace(m, [a])
+
+        results = traced(a)
+        references = m(a)
+        for res, ref in zip(results, references):
+            self.assertTrue(
+                torch.equal(ref, res), f"Tensors at don't match {ref=} {res=}"
+            )
+
     def test_all_acc_ops_registered(self):
         self.assertEqual(
             acc_normalizer._acc_ops,
@@ -2453,5 +2490,6 @@ def test_all_acc_ops_registered(self):
                 acc_ops.dtype,
                 acc_ops.isinf,
                 acc_ops.any,
+                acc_ops.tensor_split,
             },
         )

tracer/acc_tracer/acc_ops.py

@@ -3,7 +3,7 @@
 import warnings
 
 import torch  # isort:skip
-from typing import Sequence, List, cast
+from typing import Iterable, Sequence, List, cast
 
 import fx2trt_oss.tracer.acc_tracer.acc_utils as acc_utils
 import torch.nn as nn
@@ -2616,3 +2616,39 @@ def interpolate(
         mode=mode,
         align_corners=align_corners,
     )
+
+
+@register_acc_op_mapping(
+    op_and_target=("call_function", torch.tensor_split),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        (("tensor_indices_or_sections", "sections", "indices"), "indices_or_sections"),
+        ("dim", "dim", this_arg_is_optional),
+    ],
+)
+@register_acc_op_mapping(
+    op_and_target=("call_method", "tensor_split"),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        (("tensor_indices_or_sections", "sections", "indices"), "indices_or_sections"),
+        ("dim", "dim", this_arg_is_optional),
+    ],
+)
+@register_acc_op
+def tensor_split(*, input, indices_or_sections, dim=0):
+    # Need to de-coalesce the indices_or_sections because tensor_split accepts
+    # one of three kwarg signatures:
+    #  * (Tensor input, Tensor tensor_indices_or_sections, int dim)
+    #  * (Tensor input, int sections, int dim)
+    #  * (Tensor input, tuple of ints indices, int dim)
+    if isinstance(indices_or_sections, torch.Tensor):
+        indices_or_sections = indices_or_sections.tolist()
+    if isinstance(indices_or_sections, int):
+        return torch.tensor_split(input, sections=indices_or_sections, dim=dim)
+    elif isinstance(indices_or_sections, Iterable):
+        return torch.tensor_split(input, indices=tuple(indices_or_sections), dim=dim)
+    else:
+        raise RuntimeError(
+            f"Expected int, Iterable or Tensor for "
+            f"indices_or_sections arg, got: {type(indices_or_sections)}"
+        )