From 3d57623c3c22ca5c19eed1d3bf27bebff5ffb17e Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Thu, 27 Mar 2025 00:28:28 -0700
Subject: [PATCH] Module API to execute with a single Tensor.

Summary: https://github.com/pytorch/executorch/issues/8363

Differential Revision: D71955238
---
 .../ExecuTorch/Exported/ExecuTorchModule.h    | 30 ++++++++++++++++++-
 .../ExecuTorch/Exported/ExecuTorchModule.mm   | 15 ++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h
index 2e3e4e374ef..402ffa644ee 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h
@@ -173,6 +173,21 @@ __attribute__((deprecated("This API is experimental.")))
                                                  error:(NSError **)error
     NS_SWIFT_NAME(execute(_:_:));
 
+/**
+ * Executes a specific method with the provided single input tensor.
+ *
+ * The method is loaded on demand if not already loaded.
+ *
+ * @param methodName A string representing the method name.
+ * @param tensor An ExecuTorchTensor object representing the input.
+ * @param error A pointer to an NSError pointer that is set if an error occurs.
+ * @return An NSArray of ExecuTorchValue objects representing the outputs, or nil in case of an error.
+ */
+- (nullable NSArray<ExecuTorchValue *> *)executeMethod:(NSString *)methodName
+                                            withTensor:(ExecuTorchTensor *)tensor
+                                                 error:(NSError **)error
+    NS_SWIFT_NAME(execute(_:_:));
+
 /**
  * Executes the "forward" method with the provided input values.
  *
@@ -210,7 +225,7 @@ __attribute__((deprecated("This API is experimental.")))
 - (nullable NSArray<ExecuTorchValue *> *)forward:(NSError **)error;
 
 /**
- * Executes the "forward" method with no inputs.
+ * Executes the "forward" method with the provided input tensors.
  *
  * This is a convenience method that calls the executeMethod with "forward" as the method name.
  *
@@ -222,6 +237,19 @@ __attribute__((deprecated("This API is experimental.")))
                                                       error:(NSError **)error
     NS_SWIFT_NAME(forward(_:));
 
+/**
+ * Executes the "forward" method with the provided single input tensor.
+ *
+ * This is a convenience method that calls the executeMethod with "forward" as the method name.
+ *
+ * @param tensor An ExecuTorchTensor object representing the input.
+ * @param error A pointer to an NSError pointer that is set if an error occurs.
+ * @return An NSArray of ExecuTorchValue objects representing the outputs, or nil in case of an error.
+ */
+- (nullable NSArray<ExecuTorchValue *> *)forwardWithTensor:(ExecuTorchTensor *)tensor
+                                                     error:(NSError **)error
+    NS_SWIFT_NAME(forward(_:));
+
 + (instancetype)new NS_UNAVAILABLE;
 - (instancetype)init NS_UNAVAILABLE;
 
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm
index 8cdf61b26e8..98f0f555991 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm
@@ -191,6 +191,14 @@ - (BOOL)isMethodLoaded:(NSString *)methodName {
                        error:error];
 }
 
+- (nullable NSArray<ExecuTorchValue *> *)executeMethod:(NSString *)methodName
+                                            withTensor:(ExecuTorchTensor *)tensor
+                                                 error:(NSError **)error {
+  return [self executeMethod:methodName
+                  withInputs:@[[ExecuTorchValue valueWithTensor:tensor]]
+                       error:error];
+}
+
 - (nullable NSArray<ExecuTorchValue *> *)forwardWithInputs:(NSArray<ExecuTorchValue *> *)values
                                                      error:(NSError **)error {
   return [self executeMethod:@"forward"
@@ -222,4 +230,11 @@ - (BOOL)isMethodLoaded:(NSString *)methodName {
                        error:error];
 }
 
+- (nullable NSArray<ExecuTorchValue *> *)forwardWithTensor:(ExecuTorchTensor *)tensor
+                                                     error:(NSError **)error {
+  return [self executeMethod:@"forward"
+                  withInputs:@[[ExecuTorchValue valueWithTensor:tensor]]
+                       error:error];
+}
+
 @end