Add the get method to Module to return a single EValue.

Differential Revision: D61170093

Pull Request resolved: #4686

Author: shoumikhin

extension/module/module.h

@@ -158,6 +158,41 @@ class Module final {
     return execute(method_name, {});
   }
 
+  /**
+   * Retrieve the output value of a specific method with the given input.
+   * Loads the program and method before execution if needed.
+   *
+   * @param[in] method_name The name of the method to execute.
+   * @param[in] input A vector of input values to be passed to the method.
+   *
+   * @returns A Result object containing either the first output value from the
+   * method or an error to indicate failure.
+   */
+  __ET_NODISCARD
+  Result<EValue> get(
+      const std::string& method_name,
+      const std::vector<EValue>& input) {
+    auto result = ET_UNWRAP(execute(method_name, input));
+    if (result.empty()) {
+      return Error::InvalidArgument;
+    }
+    return result[0];
+  }
+
+  /**
+   * Retrieve the output value of a specific method without any input values.
+   * Loads the program and method before execution if needed.
+   *
+   * @param[in] method_name The name of the method to execute.
+   *
+   * @returns A Result object containing either the first output value from the
+   * method or an error to indicate failure.
+   */
+  __ET_NODISCARD
+  Result<EValue> get(const std::string& method_name) {
+    return get(method_name, {});
+  }
+
   /**
    * Execute the 'forward' method with the given input and retrieve output.
    * Loads the program and method before executing if needed.

extension/module/test/module_test.cpp

@@ -202,6 +202,21 @@ TEST_F(ModuleTest, TestExecuteOnCurrupted) {
   EXPECT_FALSE(result.ok());
 }
 
+TEST_F(ModuleTest, TestGet) {
+  Module module(std::getenv("RESOURCES_PATH") + std::string("/model.pte"));
+
+  std::array<float, 2> input{1, 2};
+  std::array<int32_t, 2> sizes{1, 2};
+  TensorImpl tensor(
+      ScalarType::Float, sizes.size(), sizes.data(), input.data());
+
+  const auto result = module.get("forward", {EValue(Tensor(&tensor))});
+
+  EXPECT_TRUE(result.ok());
+  const auto data = result->toTensor().const_data_ptr<float>();
+  EXPECT_NEAR(data[0], 1.5, 1e-5);
+}
+
 TEST_F(ModuleTest, TestForward) {
   auto module = std::make_unique<Module>(
       std::getenv("RESOURCES_PATH") + std::string("/model.pte"));