feat: TS Add converter support for aten::flip

core/conversion/converters/impl/select.cpp

@@ -832,6 +832,58 @@ auto select_registrations TORCHTRT_UNUSED =
                auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], layer->getOutput(0));
                LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
                return true;
+             }})
+        .pattern(
+            {"aten::flip(Tensor self, int[] dims) -> Tensor",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               auto dims = args[1].unwrapToIntList().vec();
+               auto ndims = self->getDimensions().nbDims;
+               for (auto& dim : dims) {
+                 dim = dim < 0 ? ndims + dim : dim;
+               }
+               auto dims_mask = std::vector(ndims, 0);
+               auto steps = std::vector(ndims, 1);
+               for (auto& dim : dims) {
+                 dims_mask[dim] = 1;
+                 steps[dim] = -1;
+               }
+               auto dims_mask_tensor = tensor_to_const(ctx, torch::tensor(dims_mask, torch::kInt32));
+               auto step_tensor = tensor_to_const(ctx, torch::tensor(steps, torch::kInt32));
+
+               auto self_shape_layer = ctx->net->addShape(*self);
+               TORCHTRT_CHECK(self_shape_layer, "Unable to create shape layer from node: " << util::node_info(n));
+               self_shape_layer->setName((util::node_info(n) + "_shape").c_str());
+               auto self_shape = self_shape_layer->getOutput(0);
+
+               // For dims we're flipping set start to size - 1
+               auto start_layer = add_elementwise(
+                   ctx,
+                   nvinfer1::ElementWiseOperation::kPROD,
+                   self_shape,
+                   dims_mask_tensor,
+                   util::node_info(n) + "_start");
+               auto start_tensor = start_layer->getOutput(0);
+               auto start_adjust = add_elementwise(
+                   ctx,
+                   nvinfer1::ElementWiseOperation::kSUB,
+                   start_tensor,
+                   dims_mask_tensor,
+                   util::node_info(n) + "_start_adjust");
+               start_tensor = start_adjust->getOutput(0);
+
+               // all args after slice are placeholders that will be replaced with dynamic tensors
+               auto slice_layer =
+                   ctx->net->addSlice(*self, self->getDimensions(), self->getDimensions(), self->getDimensions());
+               TORCHTRT_CHECK(slice_layer, "Unable to create slice layer from node: " << util::node_info(n));
+               slice_layer->setName((util::node_info(n) + "_slice").c_str());
+               slice_layer->setInput(1, *start_tensor);
+               slice_layer->setInput(2, *self_shape);
+               slice_layer->setInput(3, *step_tensor);
+               auto slice_out = slice_layer->getOutput(0);
+               auto out = ctx->AssociateValueAndTensor(n->outputs()[0], slice_out);
+               LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+               return true;
              }});
 
 } // namespace

tests/core/conversion/converters/test_slice.cpp

@@ -329,4 +329,32 @@ TEST(Converters, ATenSliceDynamic2ConvertsCorrectly) {
   auto trt = trt_results[0].reshape(jit_results[0].sizes());
 
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6));
-}
+}
+
+TEST(Converters, ATenFlipConvertsCorrectly) {
+  const auto graph = R"IR(
+        graph(%x.1 : Tensor):
+              %2 : int[] = prim::Constant[value=[0]]()
+              %3 : int[] = prim::Constant[value=[1]]()
+              %4 : int[] = prim::Constant[value=[0, 1]]()
+              %5 : Tensor = aten::flip(%x.1, %2)
+              %6 : Tensor = aten::flip(%x.1, %3)
+              %7 : Tensor = aten::flip(%x.1, %4)
+              return (%5, %6, %7))IR";
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::arange(8, {at::kCUDA}).to(at::kInt).reshape({2, 2, 2});
+
+  auto jit_in = at::clone(in);
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in});
+
+  auto trt_in = at::clone(in);
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {trt_in}, true);
+
+  for (size_t i = 0; i < jit_results.size(); i++) {
+    ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt_results[i], 2e-6));
+  }
+}