chore: Update generateRandomTensors uses

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/compiler.cpp

@@ -322,7 +322,7 @@ void MapInputsAndDetermineDTypes(
         est_type_opt = first_use_type_map.find(in)->second;
       }
       // traverse elements in est_type_out and spec
-      for (int i = 0; i < est_type_opt.size(); i++) {
+      for (size_t i = 0; i < est_type_opt.size(); i++) {
         if (est_type_opt[i] && !spec[i].dtype_is_user_defined) {
           // If we can calculate the type from the graph and the type was not defined by the user then use the calculated
           // type

core/conversion/evaluators/aten.cpp

@@ -264,6 +264,21 @@ auto aten_registrations TORCHTRT_UNUSED =
              },
              EvalOptions().validSchemas(
                  {"aten::size(Tensor self) -> (int[])", "aten::size.int(Tensor self, int dim) -> (int)"})})
+        .evaluator({c10::Symbol::fromQualString("aten::__getitem__"),
+                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+                      auto list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+                      auto idx = args.at(n->input(1)).unwrapToInt();
+
+                      const int64_t list_size = list.size();
+                      const int64_t normalized_idx = normalizeIndex(idx, list_size);
+                      TORCHTRT_CHECK(
+                          normalized_idx >= 0 || normalized_idx < list_size,
+                          "List index out of range (aten::__getitem__)");
+                      return list.get(normalized_idx);
+                    },
+                    EvalOptions().validSchemas({
+                        "aten::__getitem__.t(t[](a) list, int idx) -> (t(*))",
+                    })})
         .evaluator({c10::Symbol::fromQualString("aten::append"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       auto list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();

core/ir/GraphInputs.cpp

@@ -5,7 +5,7 @@ namespace torch_tensorrt {
 namespace core {
 namespace ir {
 
-void flatten_dfs(std::vector<torch_tensorrt::core::ir::Input>& flattened_inputs, std::vector<std::vector<torch_tensorrt::core::ir::Input>>& collection_inputs, 
+void flatten_dfs(std::vector<torch_tensorrt::core::ir::Input>& flattened_inputs, std::vector<std::vector<torch_tensorrt::core::ir::Input>>& collection_inputs,
                  torch::jit::IValue input_ivalue, int level, int index) {
     if (input_ivalue.isTuple()) {
       auto input_tuple = input_ivalue.toTuple();
@@ -53,7 +53,7 @@ GraphInputs::GraphInputs(std::vector<ir::Input> inputs_) {
     LOG_DEBUG("Construct GraphInput with ir::Input");
     inputs = inputs_;
     collection_inputs.resize(inputs_.size());
-    for (int i = 0; i < inputs_.size(); i++) {
+    for (size_t i = 0; i < inputs_.size(); i++) {
         collection_inputs[i].push_back(inputs_[i]);
     }
 }

core/partitioning/partitioning.cpp

@@ -400,15 +400,9 @@ PartitionedGraph segment_graph(
     if (n->kind() == torch::jit::prim::Constant) {
       continue;
     }
-<<<<<<< HEAD
-
-    if (check_node_fallback(n, global_fallback_nodes)) {
-      in_prog_trt_blk_nodes.push_back(n);
-=======
     // the outputs of trt subgraph shouldn't be collections
     if (should_run_in_trt(n, forced_fallback_ops) && !(in_prog_trt_blk_nodes.size() == 0 && is_collection(n))) {
       in_prog_trt_blk_nodes.insert(in_prog_trt_blk_nodes.begin(), n);
->>>>>>> feat: support for grouped inputs
 
       // If there is an active PyTorch block and we have passed the threshold for a valid TRT
       // block then segment and reset the active PyTorch block

core/partitioning/shape_analysis.cpp

@@ -39,15 +39,15 @@ std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomI
       std::vector<torch::jit::IValue> list;
       c10::TypePtr elementType = c10::TensorType::get();
       auto generic_list = c10::impl::GenericList(elementType);
-      for (int i = 0; i < input.second.size(); i++) {
+      for (size_t i = 0; i < input.second.size(); i++) {
         auto in = generateSingleInput(input.second[i], types[input.first][i]);
         generic_list.push_back(in.clone());
       }
       ivalue_map[input.first] = c10::IValue(generic_list);
     } else if (input.first->type()->kind() == torch::jit::TypeKind::TupleType) {
       // create tuple
       std::vector<torch::jit::IValue> list;
-      for (int i = 0; i < input.second.size(); i++) {
+      for (size_t i = 0; i < input.second.size(); i++) {
         auto in = generateSingleInput(input.second[i], types[input.first][i]);
         list.push_back(in.clone());
       }
@@ -56,7 +56,7 @@ std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomI
     } else {
       auto in = generateSingleInput(input.second[0], types[input.first][0]);
       ivalue_map[input.first] = in.clone();
-      
+
     }
   }
   return ivalue_map;

tests/core/partitioning/test_resolve_nontensor_inputs.cpp

@@ -116,7 +116,7 @@ TEST(Partitioning, ResolveNonTensorInputsCorrectly) {
   inputs.push_back(torch_tensorrt::core::ir::Input({16, 3, 3, 3}));
   inputs.push_back(torch_tensorrt::core::ir::Input({16}));
 
-  std::unordered_map<const torch::jit::Value*, std::vector<torch_tensorrt::core::ir::Input>> inputs_map;
+  torch_tensorrt::core::ir::CollectionInputSpecMap inputs_map;
   std::unordered_map<const torch::jit::Value*, std::vector<c10::optional<at::ScalarType>>> input_types;
   for (size_t i = 0; i < g->inputs().size(); ++i) {
     inputs_map.insert({g->inputs()[i], {inputs[i]}});
@@ -367,11 +367,11 @@ TEST(Partitioning, ResolveOnlyNeccessaryNonTensorInputs) {
   inputs.push_back(torch_tensorrt::core::ir::Input({4, 4}));
   inputs.push_back(torch_tensorrt::core::ir::Input({4, 4}));
 
-  std::unordered_map<const torch::jit::Value*, torch_tensorrt::core::ir::Input> inputs_map;
-  std::unordered_map<const torch::jit::Value*, c10::optional<at::ScalarType>> input_types;
+  torch_tensorrt::core::ir::CollectionInputSpecMap inputs_map;
+  std::unordered_map<const torch::jit::Value*, std::vector<c10::optional<at::ScalarType>>> input_types;
   for (size_t i = 0; i < g->inputs().size(); ++i) {
-    inputs_map.insert({g->inputs()[i], inputs[i]});
-    input_types.insert({g->inputs()[i], {at::kFloat}});
+    inputs_map.insert({g->inputs()[i], {inputs[i]}});
+    input_types.insert({g->inputs()[i], {{at::kFloat}}});
   }
   auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
   std::unordered_map<torch::jit::Node*, int> fallback_nodes;

tests/modules/hub.py

@@ -128,10 +128,10 @@
         "model": cm.ListInputTupleOutput(),
         "path": "script"
     },
-    #"bert_base_uncased": {
-    #    "model": cm.BertModule(),
-    #    "path": "trace"
-    #}
+    "bert_base_uncased": {
+        "model": cm.BertModule(),
+        "path": "trace"
+    }
 }