fix: fix failed test cases caused by partition API changes (#1460)

Signed-off-by: Bo Wang <[email protected]>

Signed-off-by: Bo Wang <[email protected]>

Author: bowang007

tests/core/partitioning/test_shape_analysis.cpp

@@ -11,7 +11,7 @@ bool checkSegmentedBlockInputShape(
   if (segmented_blocks.size() != in_shape.size())
     return false;
   for (size_t i = 0; i < segmented_blocks.size(); ++i) {
-    auto cur_block_in_shapes = segmented_blocks[i].in_shapes();
+    auto cur_block_in_shapes = segmented_blocks[i].in_opt_shapes();
     if (cur_block_in_shapes.size() != in_shape[i].size())
       return false;
     for (size_t j = 0; j < cur_block_in_shapes.size(); ++j) {

tests/core/partitioning/test_type_auto_conversion.cpp

@@ -51,10 +51,11 @@ TEST(Partitioning, ExplicitNodeAutoConversionCorrectly) {
   inputs_map.insert({g->inputs()[1], {inputs[1]}});
   input_types.insert({g->inputs()[1], {{at::kInt}}});
 
-  auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
-
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
-  torch_tensorrt::core::partitioning::partition(&ctx, input_ivalues_map);
+  ctx.input_types_map = input_types;
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
+  torch_tensorrt::core::partitioning::partition(&ctx);
   auto segmented_blocks = ctx.partitioned_blocks.begin()->second;
 
   for (auto& seg_block : segmented_blocks) {
@@ -93,10 +94,12 @@ TEST(Partitioning, ImplicitAutoConversionCorrectly) {
   inputs_map.insert({g->inputs()[0], {inputs[0]}});
   input_types.insert({g->inputs()[0], {{at::kFloat}}});
 
-  auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types);
-
+  partitioning_info.collection_input_spec_map = inputs_map;
   torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info);
-  torch_tensorrt::core::partitioning::partition(&ctx, input_ivalues_map);
+  ctx.input_types_map = input_types;
+
+  torch_tensorrt::core::partitioning::populateInputIValues(&ctx);
+  torch_tensorrt::core::partitioning::partition(&ctx);
   auto segmented_blocks = ctx.partitioned_blocks.begin()->second;
 
   for (auto& seg_block : segmented_blocks) {