Fix: fix the bug that uninitialized tensor cannot be found

This PR fixes the bug that uninitialized tensor cannot be found.
Specifically, uninitialized value is copied between subgraphs
previously, this introduces unnecessary complexity and might incur
bugs. In this PR, this uninitialized value is handled like constants:
Create one uninitialized value for each subgraph in need.

Author: bowang007

core/partitioning/partitioningctx/PartitioningCtx.cpp

@@ -15,7 +15,8 @@ PartitioningCtx::PartitioningCtx(torch::jit::Block* b, PartitioningInfo info)
 }
 
 void PartitioningCtx::_load_nodes_into_decision_map(torch::jit::Block* b) {
-  if (b->owningNode() && b->owningNode()->kind() == torch::jit::prim::Loop)
+  // won't load nodes if these nodes are in prim::loop or if these nodes are 2-level nested
+  if (b->owningNode() && (b->owningNode()->kind() == torch::jit::prim::Loop || b->owningNode()->owningBlock()->owningNode()))
     return;
 
   original_blocks.push_back(b);

core/partitioning/segmentedblock/SegmentedBlock.cpp

@@ -46,6 +46,12 @@ torch::jit::Value* SegmentedBlock::getOrAddInputForValue(torch::jit::Value* old_
       old_to_new_[old_value] = new_const->output();
       return new_const->output();
     }
+    if (node->kind() == torch::jit::prim::Uninitialized) {
+      auto new_uninitialized = g_->createUninitialized(old_value->type());
+      g_->block()->prependNode(new_uninitialized);
+      old_to_new_[old_value] = new_uninitialized->output();
+      return new_uninitialized->output();
+    }
     auto new_value = g_->block()->addInput();
     // every time when we addInput, we push back the corresponding lowering graph torch::jit::Value to our raw_inputs
     inputs_.push_back(old_value);