Update docs on Module new APIs.

docs/source/extension-module.md

@@ -2,7 +2,7 @@
 
 **Author:** [Anthony Shoumikhin](https://github.com/shoumikhin)
 
-In the [Running an ExecuTorch Model in C++ Tutorial](running-a-model-cpp-tutorial.md), we explored the lower-level ExecuTorch APIs for running an exported model. While these APIs offer zero overhead, great flexibility, and control, they can be verbose and complex for regular use. To simplify this and resemble PyTorch's eager mode in Python, we introduce the Module facade APIs over the regular ExecuTorch runtime APIs. The Module APIs provide the same flexibility but default to commonly used components like `DataLoader` and `MemoryAllocator`, hiding most intricate details.
+In the [Running an ExecuTorch Model in C++ Tutorial](running-a-model-cpp-tutorial.md), we explored the lower-level ExecuTorch APIs for running an exported model. While these APIs offer zero overhead, great flexibility, and control, they can be verbose and complex for regular use. To simplify this and resemble PyTorch's eager mode in Python, we introduce the `Module` facade APIs over the regular ExecuTorch runtime APIs. The `Module` APIs provide the same flexibility but default to commonly used components like `DataLoader` and `MemoryAllocator`, hiding most intricate details.
 
 ## Example
 
@@ -37,7 +37,7 @@ The code now boils down to creating a `Module` and calling `forward()` on it, wi
 
 ### Creating a Module
 
-Creating a `Module` object is an extremely fast operation that does not involve significant processing time or memory allocation. The actual loading of a `Program` and a `Method` happens lazily on the first inference unless explicitly requested with a dedicated API.
+Creating a `Module` object is a fast operation that does not involve significant processing time or memory allocation. The actual loading of a `Program` and a `Method` happens lazily on the first inference unless explicitly requested with a dedicated API.
 
 ```cpp
 Module module("/path/to/model.pte");
@@ -60,31 +60,32 @@ const auto error = module.load_method("forward");
 
 assert(module.is_method_loaded("forward"));
 ```
-Note: the `Program` is loaded automatically before any `Method` is loaded. Subsequent attemps to load them have no effect if one of the previous attemps was successful.
 
-You can also force-load the "forward" method with a convenience syntax:
+You can also use the convenience function to load the `forward` method:
 
 ```cpp
 const auto error = module.load_forward();
 
 assert(module.is_method_loaded("forward"));
 ```
 
+**Note:** The `Program` is loaded automatically before any `Method` is loaded. Subsequent attempts to load them have no effect if a previous attempt was successful.
+
 ### Querying for Metadata
 
-Get a set of method names that a Module contains udsing the `method_names()` function:
+Get a set of method names that a `Module` contains using the `method_names()` function:
 
 ```cpp
 const auto method_names = module.method_names();
 
 if (method_names.ok()) {
-  assert(method_names.count("forward"));
+  assert(method_names->count("forward"));
 }
 ```
 
-Note: `method_names()` will try to force-load the `Program` when called the first time.
+**Note:** `method_names()` will force-load the `Program` when called for the first time.
 
-Introspect miscellaneous metadata about a particular method via `MethodMeta` struct returned by `method_meta()` function:
+To introspect miscellaneous metadata about a particular method, use the `method_meta()` function, which returns a `MethodMeta` struct:
 
 ```cpp
 const auto method_meta = module.method_meta("forward");
@@ -94,47 +95,123 @@ if (method_meta.ok()) {
   assert(method_meta->num_inputs() > 1);
 
   const auto input_meta = method_meta->input_tensor_meta(0);
-
   if (input_meta.ok()) {
     assert(input_meta->scalar_type() == ScalarType::Float);
   }
-  const auto output_meta = meta->output_tensor_meta(0);
 
+  const auto output_meta = method_meta->output_tensor_meta(0);
   if (output_meta.ok()) {
     assert(output_meta->sizes().size() == 1);
   }
 }
 ```
 
-Note: `method_meta()` will try to force-load the `Method` when called for the first time.
+**Note:** `method_meta()` will also force-load the `Method` the first time it is called.
 
-### Perform an Inference
+### Performing an Inference
 
-Assuming that the `Program`'s method names and their input format is known ahead of time, we rarely need to query for those and can run the methods directly by name using the `execute()` function:
+Assuming the `Program`'s method names and their input format are known ahead of time, you can run methods directly by name using the `execute()` function:
 
 ```cpp
-const auto result = module.execute("forward", Tensor(&tensor));
+const auto result = module.execute("forward", tensor);
+```
+
+For the standard `forward()` method, the above can be simplified:
+
+```cpp
+const auto result = module.forward(tensor);
 ```
 
-Which can also be simplified for the standard `forward()` method name as:
+**Note:** `execute()` or `forward()` will load the `Program` and the `Method` the first time they are called. Therefore, the first inference will take longer, as the model is loaded lazily and prepared for execution unless it was explicitly loaded earlier.
+
+### Setting Input and Output
+
+You can set individual input and output values for methods with the following APIs.
+
+#### Setting Inputs
+
+Inputs can be any `EValue`, which includes tensors, scalars, lists, and other supported types. To set a specific input value for a method:
+
+```cpp
+module.set_input("forward", input_value, input_index);
+```
+
+- `input_value` is an `EValue` representing the input you want to set.
+- `input_index` is the zero-based index of the input to set.
+
+For example, to set the first input tensor:
+
+```cpp
+module.set_input("forward", tensor_value, 0);
+```
+
+You can also set multiple inputs at once:
+
+```cpp
+std::vector<runtime::EValue> inputs = {input1, input2, input3};
+module.set_inputs("forward", inputs);
+```
+
+**Note:** You can skip the method name argument for the `forward()` method.
+
+By pre-setting all inputs, you can perform an inference without passing any arguments:
 
 ```cpp
-const auto result = module.forward(Tensor(&tensor));
+const auto result = module.forward();
 ```
 
-Note: `execute()` or `forward()` will try to force load the `Program` and the `Method` when called for the first time. Therefore, the first inference will take more time than subsequent ones as it loads the model lazily and prepares it for execution unless the `Program` or `Method` was loaded explicitly earlier using the corresponding functions.
+Or just setting and then passing the inputs partially:
+
+```cpp
+// Set the second input ahead of time.
+module.set_input(input_value_1, 1);
+
+// Execute the method, providing the first input at call time.
+const auto result = module.forward(input_value_0);
+```
+
+**Note:** The pre-set inputs are stored in the `Module` and can be reused multiple times for the next executions.
+
+Don't forget to clear or reset the inputs if you don't need them anymore by setting them to default-constructed `EValue`:
+
+```cpp
+module.set_input(runtime::EValue(), 1);
+```
+
+#### Setting Outputs
+
+Only outputs of type Tensor can be set at runtime, and they must not be memory-planned at model export time. Memory-planned tensors are preallocated during model export and cannot be replaced.
+
+To set the output tensor for a specific method:
+
+```cpp
+module.set_output("forward", output_tensor, output_index);
+```
+
+- `output_tensor` is an `EValue` containing the tensor you want to set as the output.
+- `output_index` is the zero-based index of the output to set.
+
+**Note:** Ensure that the output tensor you're setting matches the expected shape and data type of the method's output.
+
+You can skip the method name for `forward()` and the index for the first output:
+
+```cpp
+module.set_output(output_tensor);
+```
+
+**Note:** The pre-set outputs are stored in the `Module` and can be reused multiple times for the next executions, just like inputs.
 
 ### Result and Error Types
 
-Most of the ExecuTorch APIs, including those described above, return either `Result` or `Error` types. Let's understand what those are:
+Most of the ExecuTorch APIs return either `Result` or `Error` types:
 
-* [`Error`](https://github.com/pytorch/executorch/blob/main/runtime/core/error.h) is a C++ enum containing a collection of valid error codes, where the default is `Error::Ok`, denoting success.
+- [`Error`](https://github.com/pytorch/executorch/blob/main/runtime/core/error.h) is a C++ enum containing valid error codes. The default is `Error::Ok`, denoting success.
 
-* [`Result`](https://github.com/pytorch/executorch/blob/main/runtime/core/result.h) can hold either an `Error` if the operation has failed or a payload, i.e., the actual result of the operation like an `EValue` wrapping a `Tensor` or any other standard C++ data type if the operation succeeded. To check if `Result` has a valid value, call the `ok()` function. To get the `Error` use the `error()` function, and to get the actual data, use the overloaded `get()` function or dereferencing pointer operators like `*` and `->`.
+- [`Result`](https://github.com/pytorch/executorch/blob/main/runtime/core/result.h) can hold either an `Error` if the operation fails, or a payload such as an `EValue` wrapping a `Tensor` if successful. To check if a `Result` is valid, call `ok()`. To retrieve the `Error`, use `error()`, and to get the data, use `get()` or dereference operators like `*` and `->`.
 
-### Profile the Module
+### Profiling the Module
 
-Use [ExecuTorch Dump](etdump.md) to trace model execution. Create an instance of the `ETDumpGen` class and pass it to the `Module` constructor. After executing a method, save the `ETDump` to a file for further analysis. You can capture multiple executions in a single trace if desired.
+Use [ExecuTorch Dump](etdump.md) to trace model execution. Create an `ETDumpGen` instance and pass it to the `Module` constructor. After executing a method, save the `ETDump` data to a file for further analysis:
 
 ```cpp
 #include <fstream>
@@ -147,7 +224,7 @@ using namespace ::executorch::extension;
 
 Module module("/path/to/model.pte", Module::LoadMode::MmapUseMlock, std::make_unique<ETDumpGen>());
 
-// Execute a method, e.g. module.forward(...); or module.execute("my_method", ...);
+// Execute a method, e.g., module.forward(...); or module.execute("my_method", ...);
 
 if (auto* etdump = dynamic_cast<ETDumpGen*>(module.event_tracer())) {
   const auto trace = etdump->get_etdump_data();
@@ -162,3 +239,7 @@ if (auto* etdump = dynamic_cast<ETDumpGen*>(module.event_tracer())) {
   }
 }
 ```
+
+# Conclusion
+
+The `Module` APIs provide a simplified interface for running ExecuTorch models in C++, closely resembling the experience of PyTorch's eager mode. By abstracting away the complexities of the lower-level runtime APIs, developers can focus on model execution without worrying about the underlying details.