fix: resolve review comments

Author: JYMiracle305

infini_train/include/nn/modules/container.h

@@ -35,12 +35,12 @@ class ModuleList : public CloneableModule<ModuleList> {
 
     std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
 
-    auto begin() { return module_list_.begin(); }
-    auto end() { return module_list_.end(); }
-    auto begin() const { return module_list_.begin(); }
-    auto end() const { return module_list_.end(); }
+    std::vector<std::shared_ptr<Module>>::iterator begin();
+    std::vector<std::shared_ptr<Module>>::iterator end();
+    std::vector<std::shared_ptr<Module>>::const_iterator begin() const;
+    std::vector<std::shared_ptr<Module>>::const_iterator end() const;
 
-    std::shared_ptr<Module> &operator[](std::size_t idx) { return module_list_.at(idx); }
+    std::shared_ptr<Module> &operator[](std::size_t idx);
 
 private:
     std::vector<std::shared_ptr<Module>> module_list_;

infini_train/include/nn/parallel/pp/pipeline_stage.h

@@ -22,19 +22,19 @@ class PipelineStage {
     std::vector<std::shared_ptr<Tensor>> ForwardOneChunk(const std::vector<std::shared_ptr<Tensor>> &inputs,
                                                          int local_chunk_idx = 0);
 
-    bool IsFirstStage() const { return stage_index_ == 0; }
-    bool IsLastStage() const { return stage_index_ == num_stages_ - 1; }
-
-    int stage_index() const { return stage_index_; }
-    int prev_rank() const { return prev_rank_; }
-    int next_rank() const { return next_rank_; }
-    int num_stages() const { return num_stages_; }
-
-    const Device *device() const { return device_; }
-    const std::vector<std::vector<int64_t>> &recv_shape() const { return recv_shape_; }
-    std::shared_ptr<Optimizer> optimizer() { return optimizer_; }
-    const auto &chunks() { return chunks_; }
-    auto *mutable_chunks() { return &chunks_; }
+    bool IsFirstStage() const;
+    bool IsLastStage() const;
+
+    int stage_index() const;
+    int prev_rank() const;
+    int next_rank() const;
+    int num_stages() const;
+
+    const Device *device() const;
+    const std::vector<std::vector<int64_t>> &recv_shape() const;
+    std::shared_ptr<Optimizer> optimizer();
+    const std::vector<std::shared_ptr<Module>> &chunks();
+    std::vector<std::shared_ptr<Module>> *mutable_chunks();
 
 private:
     int stage_index_ = -1;

infini_train/src/nn/modules/container.cc

@@ -41,4 +41,12 @@ ModuleList::ModuleList(std::vector<std::shared_ptr<Module>> &&layers)
 std::vector<std::shared_ptr<Tensor>> ModuleList::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
     LOG(FATAL) << "Not implemented";
 }
+
+std::vector<std::shared_ptr<Module>>::iterator ModuleList::begin() { return module_list_.begin(); }
+std::vector<std::shared_ptr<Module>>::iterator ModuleList::end() { return module_list_.end(); }
+std::vector<std::shared_ptr<Module>>::const_iterator ModuleList::begin() const { return module_list_.begin(); }
+std::vector<std::shared_ptr<Module>>::const_iterator ModuleList::end() const { return module_list_.end(); }
+
+std::shared_ptr<Module> &ModuleList::operator[](std::size_t idx) { return module_list_.at(idx); }
+
 } // namespace infini_train::nn

infini_train/src/nn/parallel/pp/pipeline_stage.cc

@@ -28,4 +28,17 @@ std::vector<std::shared_ptr<Tensor>> PipelineStage::ForwardOneChunk(const std::v
     return chunks_[local_chunk_idx]->Forward(inputs);
 }
 
+bool PipelineStage::IsFirstStage() const { return stage_index_ == 0; }
+bool PipelineStage::IsLastStage() const { return stage_index_ == num_stages_ - 1; }
+
+int PipelineStage::stage_index() const { return stage_index_; }
+int PipelineStage::prev_rank() const { return prev_rank_; }
+int PipelineStage::next_rank() const { return next_rank_; }
+int PipelineStage::num_stages() const { return num_stages_; }
+
+const Device *PipelineStage::device() const { return device_; }
+const std::vector<std::vector<int64_t>> &PipelineStage::recv_shape() const { return recv_shape_; }
+std::shared_ptr<Optimizer> PipelineStage::optimizer() { return optimizer_; }
+const std::vector<std::shared_ptr<Module>> &PipelineStage::chunks() { return chunks_; }
+std::vector<std::shared_ptr<Module>> *PipelineStage::mutable_chunks() { return &chunks_; }
 } // namespace infini_train::nn::parallel