Work around torch.jit.trace bug upon Gather-ScatterElements sequence

## Problems
Gather was added to SKIP_LIST_FOR_SUBGRAPH_TRACE in 7563948, but it turned out that it only creates a new problem without solving any existing problems.

Two symptoms have been identified
1. Putting `Gather` in `SKIP_LIST_FOR_SUBGRAPH_TRACE` doesn't help torch.jit.trace capture the input order correctly.
2. For some reason, putting `Gather` in `SKIP_LIST_FOR_SUBGRAPH_TRACE` causes another type of mis-tracing upon Gather-ScatterElements sequence. It thinks ScatterElements and Gather takes the same input, when in fact ScatterElements takes Gather's output as input.

## Main Changes
* Partially reverted 7563948 -- removed Gather from `SKIP_LIST_FOR_SUBGRAPH_TRACE`. This fixes problem 2
* Added explicit unit tests for both problem 1 and 2

---------

Signed-off-by: Kyunggeun Lee <kyunggeu@qti.qualcomm.com>

Author: kyunggeunlee

TrainingExtensions/torch/src/python/aimet_torch/meta/connectedgraph.py

@@ -108,7 +108,6 @@
     aimet_modules.DepthToSpaceCRDMode,
     aimet_modules.Addmm,
     aimet_modules.Baddbmm,
-    aimet_modules.Gather,
     aimet_modules.GridSample,
 ]
 

TrainingExtensions/torch/test/python/test_connectedgraph.py

@@ -1428,6 +1428,11 @@ def forward(self, x, y):
         (aimet_modules.Equal, (torch.randn(3),), torch.randn(3)),
         (aimet_modules.FloorDivide, (torch.randn(3),), torch.randn(3)),
         (aimet_modules.Fmod, (torch.randn(3),), torch.randn(3)),
+        (
+            aimet_modules.Gather,
+            (torch.randn(3, 3), torch.tensor(1, dtype=torch.int64)),
+            torch.tensor([[0, 0], [1, 1]]),
+        ),
         (aimet_modules.Greater, (torch.randn(3),), torch.randn(3)),
         (aimet_modules.GreaterEqual, (torch.randn(3),), torch.randn(3)),
         (aimet_modules.GridSample, (torch.randn(1, 1, 2, 2),), torch.randn(1, 2, 2, 2)),
@@ -1454,6 +1459,9 @@ def forward(self, x, y):
     ],
 )
 def test_nary_operator_input_ordering(layer, inputs, buffer):
+    if issubclass(layer, aimet_modules.Gather):
+        pytest.xfail(reason="Known failure. FIX ASAP")
+
     class Model(torch.nn.Module):
         def __init__(self, layer):
             super(Model, self).__init__()
@@ -1468,3 +1476,23 @@ def forward(self, *x):
     assert sim.connected_graph.ordered_ops[0].inputs[-1].is_const
     for inp in sim.connected_graph.ordered_ops[0].inputs[:-1]:
         assert inp.is_model_input
+
+
+def test_gather_scatter_element_sequence():
+    class GatherScatter(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.gather = custom.Gather()
+            self.scatter_elements = custom.ScatterElements(dim=3)
+
+        def forward(self, input, indices):
+            gather_out = self.gather(input, 3, indices)
+            return self.scatter_elements(input, indices, gather_out)
+
+    input = torch.randn(1, 1, 3, 128)
+    indices = torch.zeros(1, 1, 3, 16, dtype=torch.int64)
+    model = GatherScatter()
+    sim = aimet_torch.QuantizationSimModel(model, (input, indices))
+
+    gather, scatter_elements = sim.connected_graph.ordered_ops
+    assert scatter_elements.inputs[-1] is gather.outputs[0]