Add support in training_module to get attributes

Summary: As title

Differential Revision: D78907316

Author: silverguo

extension/module/module.cpp

@@ -230,11 +230,7 @@ runtime::Error Module::load_method(
 
 ET_NODISCARD runtime::Result<Method*> Module::method(
     const std::string& method_name) {
-  ET_CHECK_OR_RETURN_ERROR(
-      methods_.count(method_name) > 0,
-      InvalidArgument,
-      "no such method in program: %s",
-      method_name.c_str());
+  ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
   return methods_[method_name].method.get();
 }
 

extension/training/module/test/training_module_test.cpp

@@ -88,6 +88,12 @@ TEST_F(TrainingModuleTest, JointGraphTest) {
   ASSERT_EQ(param.find("linear.weight")->second.dim(), 2);
   ASSERT_EQ(param.find("linear.bias")->second.sizes()[0], 3);
   ASSERT_EQ(param.find("linear.bias")->second.dim(), 1);
+
+  // Test attributes for pte only model
+  auto attributes_res = mod.named_attributes("forward");
+  ASSERT_EQ(attributes_res.error(), Error::Ok);
+  auto& attributes = attributes_res.get();
+  ASSERT_EQ(attributes.size(), 0);
 }
 
 TEST_F(TrainingModuleTest, NonTrainingModuleTest) {
@@ -152,4 +158,18 @@ TEST_F(TrainingModuleTest, SeperateDataTest) {
   auto res = mod.execute_forward_backward("forward", inputs);
   ASSERT_EQ(res.error(), Error::Ok);
   ASSERT_EQ(res.get().size(), 1);
+
+  // Test Attributes for pte + ptd model
+  auto attributes_res = mod.named_attributes("forward");
+  ASSERT_EQ(attributes_res.error(), Error::Ok);
+  auto& attributes = attributes_res.get();
+  ASSERT_EQ(attributes.size(), 2);
+  ASSERT_NE(attributes.find("linear.weight"), attributes.end());
+  ASSERT_NE(attributes.find("linear.bias"), attributes.end());
+
+  ASSERT_EQ(attributes.find("linear.weight")->second.sizes()[0], 3);
+  ASSERT_EQ(attributes.find("linear.weight")->second.sizes()[1], 3);
+  ASSERT_EQ(attributes.find("linear.weight")->second.dim(), 2);
+  ASSERT_EQ(attributes.find("linear.bias")->second.sizes()[0], 3);
+  ASSERT_EQ(attributes.find("linear.bias")->second.dim(), 1);
 }

extension/training/module/training_module.cpp

@@ -154,6 +154,41 @@ TrainingModule::named_gradients(const std::string& method_name) {
   return method_named_gradients_.at(method_name);
 }
 
+runtime::Result<const std::map<std::string_view, executorch::aten::Tensor>>
+TrainingModule::named_attributes(const std::string& method_name) {
+  // If we haven't seen this method before, populate the dict.
+  if (method_named_attributes_.find(method_name) ==
+      method_named_attributes_.end()) {
+    method_named_attributes_.insert({method_name, {}});
+
+    // get method metadata
+    auto meta_res = executorch::extension::Module::method_meta(method_name);
+    if (!meta_res.ok()) {
+      return meta_res.error();
+    }
+    // get method
+    auto method_res = executorch::extension::Module::method(method_name);
+    if (!method_res.ok()) {
+      return method_res.error();
+    }
+    // get tensor by name
+    for (int idx = 0; idx < meta_res->num_attributes(); idx++) {
+      const auto tensor_res = meta_res->attribute_tensor_meta(idx);
+      if (!tensor_res.ok()) {
+        return tensor_res.error();
+      }
+      const auto tensorName = tensor_res.get().name();
+      const auto attribute_res = (*method_res)->get_attribute(tensorName);
+      if (!attribute_res.ok()) {
+        return attribute_res.error();
+      }
+      method_named_attributes_.at(method_name)
+          .insert({tensorName, attribute_res.get()});
+    }
+  }
+  return method_named_attributes_.at(method_name);
+}
+
 } // namespace training
 } // namespace extension
 } // namespace executorch

extension/training/module/training_module.h

@@ -93,6 +93,19 @@ class ET_EXPERIMENTAL TrainingModule final
   runtime::Result<const std::map<std::string_view, executorch::aten::Tensor>>
   named_gradients(const std::string& method_name);
 
+  /**
+   * Retrieve the attributes for a method.
+   *
+   * @param[in] method_name The name of the  method to get the
+   * attributes for.
+   *
+   * @returns A Result object containing a map of the fully qualified name to
+   * attribute tensor.
+   */
+  ET_EXPERIMENTAL
+  runtime::Result<const std::map<std::string_view, executorch::aten::Tensor>>
+  named_attributes(const std::string& method_name);
+
  private:
   std::unordered_map<
       std::string,
@@ -103,6 +116,11 @@ class ET_EXPERIMENTAL TrainingModule final
       std::string,
       std::map<std::string_view, executorch::aten::Tensor>>
       method_named_parameters_;
+
+  std::unordered_map<
+      std::string,
+      std::map<std::string_view, executorch::aten::Tensor>>
+      method_named_attributes_;
 };
 
 } // namespace training