Repository: deepjavalibrary/djl
Pull request #2234 (2022-12-16 04:46:19)

[PyTorch] use runMethod to replace forward function
## Description ##

Support run different methods on pytorch with different names


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2234?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **71.54**% // Decreases project coverage by **`-0.54%`** :warning:
> Coverage data is based on head [(`1c6de63`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2234?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 73.70% of modified lines in pull request are covered.

> :exclamation: Current head 1c6de63 differs from pull request most recent head fab1dac. Consider uploading reports for the commit fab1dac to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2234      +/-   ##
============================================
- Coverage     72.08%   71.54%   -0.55%     
- Complexity     5126     6343    +1217     
============================================
  Files           473      631     +158     
  Lines         21970    28079    +6109     
  Branches       2351     2989     +638     
============================================
+ Hits          15838    20090    +4252     
- Misses         4925     6526    +1601     
- Partials       1207     1463     +256     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2234?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...ain/java/ai/djl/modality/cv/util/NDImageUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3V0aWwvTkRJbWFnZVV0aWxzLmphdmE=) | `59.21% <0.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/nlp/Decoder.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9EZWNvZGVyLmphdmE=) | `63.63% <ø> (ø)` | |
| ... and [564 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2234/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2234?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/docs/pytorch/how_to_convert_your_model_to_torchscript.md b/docs/pytorch/how_to_convert_your_model_to_torchscript.md
index 6b64559bb..7387884d1 100644
--- a/docs/pytorch/how_to_convert_your_model_to_torchscript.md
+++ b/docs/pytorch/how_to_convert_your_model_to_torchscript.md
@@ -7,6 +7,9 @@ This works best when your model doesn't have control flow.
 If you do have control flow, you will need to use the scripting approach.
 In DJL, we use tracing to create TorchScript for our ModelZoo models.
 
+
+## Basic conversion
+
 Here is an example of tracing in actions:
 
 ```python
@@ -29,4 +32,43 @@ traced_script_module = torch.jit.trace(model, example)
 traced_script_module.save("traced_resnet_model.pt")
 ```
 
+Then you can just use the saved model in DJL like all other models.
+
+## Advanced case
+
+In some cases, you may have a method name that is not `forward` in pytorch (HuggingFace) like one below:
+
+```python
+from transformers import CLIPModel, CLIPProcessor
+import torch
+from PIL import Image
+import requests
+model_name = "openai/clip-vit-base-patch32"
+
+model = CLIPModel.from_pretrained(model_name, torchscript=True, return_dict=False)
+processor = CLIPProcessor.from_pretrained(model_name)
+
+test_text = "this is a cat"
+text_inputs = processor(text=test_text, return_tensors="pt")
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+image_inputs = processor(images=image, return_tensors="pt")
+
+converted = torch.jit.trace_module(model,  {'get_text_features': [text_inputs['input_ids'], text_inputs['attention_mask']],
+                                            'get_image_features': [image_inputs['pixel_values']]})
+
+torch.jit.save(converted, "cliptext.pt")
+```
+
+You can trace by using the `torch.traceModule` function.
+
+To run inference with such model in DJL, you could provide a placeholder NDArray like below:
+
+```
+NDArray array = NDManager.create("");
+array.setName("module_method:get_text_features");
+```
+
+to tell which entryPoint method you would like to use. The placeholder NDArray will NOT be used in inference.
+
 For more information on exporting TorchScript, see the [Loading a torch script model in C++ tutorial](https://pytorch.org/tutorials/advanced/cpp_export.html).
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java
index 0a1da4b88..69e213bcf 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java
@@ -31,6 +31,9 @@ import java.util.regex.Pattern;
 /** IValueUtils is utility class to deal with IValue in PyTorch. */
 public final class IValueUtils {
 
+    private static final Pattern PATTERN_LIST = Pattern.compile("\\w+\\[]");
+    private static final Pattern PATTERN_TUPLE = Pattern.compile("\\w+\\(\\)");
+
     private IValueUtils() {}
 
     /**
@@ -42,11 +45,15 @@ public final class IValueUtils {
      * @return the result {@link NDList}
      */
     public static NDList forward(PtSymbolBlock block, NDList inputs, boolean isTrain) {
-        IValue[] iValues = getInputs(inputs);
-        long[] iValueHandles = Arrays.stream(iValues).mapToLong(IValue::getHandle).toArray();
-        long result = PyTorchLibrary.LIB.moduleForward(block.getHandle(), iValueHandles, isTrain);
+        Pair<IValue[], String> inputPair = getInputs(inputs);
+        IValue[] ivalues = inputPair.getKey();
+        String method = inputPair.getValue();
+        long[] iValueHandles = Arrays.stream(ivalues).mapToLong(IValue::getHandle).toArray();
+        long result =
+                PyTorchLibrary.LIB.moduleRunMethod(
+                        block.getHandle(), method, iValueHandles, isTrain);
         PtNDManager manager = (PtNDManager) inputs.get(0).getManager();
-        Arrays.stream(iValues).forEach(IValue::close);
+        Arrays.stream(ivalues).forEach(IValue::close);
         try (IValue iValue = new IValue(result)) {
             return iValue.toNDList(manager);
         }
@@ -60,8 +67,21 @@ public final class IValueUtils {
      * @return the result {@link IValue}
      */
     public static IValue forward(PtSymbolBlock block, IValue... inputs) {
+        return runMethod(block, "forward", inputs);
+    }
+
+    /**
+     * Runs the method of PyTorch module.
+     *
+     * @param block the block that contains PyTorch module
+     * @param methodName the name of method for calling
+     * @param inputs the input {@link IValue}
+     * @return the result {@link IValue}
+     */
+    public static IValue runMethod(PtSymbolBlock block, String methodName, IValue... inputs) {
         long[] handles = Arrays.stream(inputs).mapToLong(IValue::getHandle).toArray();
-        return new IValue(PyTorchLibrary.LIB.moduleForward(block.getHandle(), handles, false));
+        return new IValue(
+                PyTorchLibrary.LIB.moduleRunMethod(block.getHandle(), methodName, handles, false));
     }
 
     private static int addToMap(
@@ -74,9 +94,10 @@ public final class IValueUtils {
                 });
     }
 
-    static IValue[] getInputs(NDList ndList) {
+    static Pair<IValue[], String> getInputs(NDList ndList) {
         List<PairList<String, PtNDArray>> outputs = new ArrayList<>();
         Map<String, Integer> indexMap = new ConcurrentHashMap<>();
+        String methodName = "forward";
         for (NDArray array : ndList) {
             String name = array.getName();
             if (name != null && name.contains(".")) {
@@ -84,11 +105,13 @@ public final class IValueUtils {
                 int index = addToMap(indexMap, strings[0], outputs);
                 PairList<String, PtNDArray> pl = outputs.get(index);
                 pl.add(strings[1], (PtNDArray) array);
-            } else if (name != null && Pattern.matches("\\w+\\[]", name)) {
+            } else if (name != null && name.startsWith("module_method:")) {
+                methodName = name.substring(14);
+            } else if (name != null && PATTERN_LIST.matcher(name).matches()) {
                 int index = addToMap(indexMap, name, outputs);
                 PairList<String, PtNDArray> pl = outputs.get(index);
                 pl.add("[]", (PtNDArray) array);
-            } else if (name != null && Pattern.matches("\\w+\\(\\)", name)) {
+            } else if (name != null && PATTERN_TUPLE.matcher(name).matches()) {
                 int index = addToMap(indexMap, name, outputs);
                 PairList<String, PtNDArray> pl = outputs.get(index);
                 pl.add("()", (PtNDArray) array);
@@ -121,6 +144,6 @@ public final class IValueUtils {
                 ret[i] = IValue.stringMapFrom(map);
             }
         }
-        return ret;
+        return new Pair<>(ret, methodName);
     }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 47fb0b724..411c4843a 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -1629,6 +1629,10 @@ public final class JniUtils {
         return list;
     }
 
+    public static String[] getMethodNames(PtSymbolBlock block) {
+        return PyTorchLibrary.LIB.moduleGetMethodNames(block.getHandle());
+    }
+
     public static void enableInferenceMode(PtSymbolBlock block) {
         PyTorchLibrary.LIB.moduleEval(block.getHandle());
     }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 53b043b4e..7196b55be 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -505,7 +505,8 @@ final class PyTorchLibrary {
 
     native void moduleTrain(long handle);
 
-    native long moduleForward(long moduleHandle, long[] iValueHandles, boolean isTrain);
+    native long moduleRunMethod(
+            long moduleHandle, String methodName, long[] iValueHandles, boolean isTrain);
 
     native void setGraphExecutorOptimize(boolean enabled);
 
@@ -515,6 +516,8 @@ final class PyTorchLibrary {
 
     native String[] moduleGetParamNames(long moduleHandle);
 
+    native String[] moduleGetMethodNames(long moduleHandle);
+
     native long iValueFromTensor(long tensorHandle);
 
     native long iValueFromBool(boolean value);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc
index 9b871f775..f103aced2 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc
@@ -210,11 +210,12 @@ JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleTrain(
   API_END()
 }
 
-JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleForward(
-    JNIEnv* env, jobject jthis, jlong module_handle, jlongArray jivalue_ptrs, jboolean jis_train) {
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleRunMethod(JNIEnv* env, jobject jthis,
+    jlong module_handle, jstring jmethod_name, jlongArray jivalue_ptrs, jboolean jis_train) {
   API_BEGIN()
   auto* module_ptr = reinterpret_cast<torch::jit::script::Module*>(module_handle);
   size_t len = env->GetArrayLength(jivalue_ptrs);
+  auto method_name = djl::utils::jni::GetStringFromJString(env, jmethod_name);
   jlong* jptrs = env->GetLongArrayElements(jivalue_ptrs, JNI_FALSE);
   std::vector<torch::IValue> inputs;
   inputs.reserve(len);
@@ -223,11 +224,11 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleForward(
   }
   torch::IValue output = [&]() {
     if (jis_train) {
-      return module_ptr->forward(inputs);
+      return module_ptr->get_method(method_name)(std::move(inputs));
     }
     // disable autograd
     JITCallGuard guard;
-    return module_ptr->forward(inputs);
+    return module_ptr->get_method(method_name)(std::move(inputs));
   }();
   env->ReleaseLongArrayElements(jivalue_ptrs, jptrs, djl::utils::jni::RELEASE_MODE);
   const auto* result_ptr = new torch::IValue(output);
@@ -281,3 +282,15 @@ JNIEXPORT jobjectArray JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleGetP
   return djl::utils::jni::GetStringArrayFromVec(env, jptrs);
   API_END_RETURN()
 }
+
+JNIEXPORT jobjectArray JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleGetMethodNames(
+    JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  auto* module_ptr = reinterpret_cast<torch::jit::script::Module*>(jhandle);
+  std::vector<std::string> jptrs;
+  for (const auto& method : module_ptr->get_methods()) {
+    jptrs.push_back(method.name());
+  }
+  return djl::utils::jni::GetStringArrayFromVec(env, jptrs);
+  API_END_RETURN()
+}


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/TorchScriptTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/TorchScriptTest.java
index 3b5899517..54b19c756 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/TorchScriptTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/TorchScriptTest.java
@@ -96,4 +96,19 @@ public class TorchScriptTest {
             }
         }
     }
+
+    @Test
+    public void testGetMethodNames() throws ModelException, IOException {
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, NDList.class)
+                        .optModelUrls("djl://ai.djl.pytorch/resnet/0.0.1/traced_resnet18")
+                        .optProgress(new ProgressBar())
+                        .build();
+        try (ZooModel<NDList, NDList> model = criteria.loadModel()) {
+            PtSymbolBlock block = (PtSymbolBlock) model.getBlock();
+            String[] result = JniUtils.getMethodNames(block);
+            Assert.assertEquals(result[0], "forward");
+        }
+    }
 }
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
index 2afbb10c2..0d4769c84 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
@@ -16,6 +16,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
+import ai.djl.util.Pair;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -35,10 +36,12 @@ public class IValueUtilsTest {
             // the NDList is mapped to (input1: Tuple(Tensor))
             input.attach(manager);
 
-            IValue[] iValues = IValueUtils.getInputs(input);
+            Pair<IValue[], String> result = IValueUtils.getInputs(input);
+            IValue[] iValues = result.getKey();
             Assert.assertEquals(iValues.length, 1);
             Assert.assertTrue(iValues[0].isTuple());
             Assert.assertEquals(iValues[0].toIValueTuple().length, 2);
+            Assert.assertEquals(result.getValue(), "forward");
 
             Arrays.stream(iValues).forEach(IValue::close);
         }
@@ -63,7 +66,7 @@ public class IValueUtilsTest {
             // the 2nd part is the key of each value in the dict
             input.attach(manager);
 
-            IValue[] iValues = IValueUtils.getInputs(input);
+            IValue[] iValues = IValueUtils.getInputs(input).getKey();
             Assert.assertEquals(iValues.length, 2);
             Assert.assertTrue(iValues[0].isMap());
             Assert.assertEquals(iValues[1].toIValueMap().size(), 3);
@@ -81,14 +84,18 @@ public class IValueUtilsTest {
             array2.setName("input1[]");
             NDArray array3 = manager.ones(new Shape(1));
             array3.setName("input2");
-            NDList input = new NDList(array1, array2, array3);
+            NDArray array4 = manager.create("Hello World");
+            array4.setName("module_method:get_text_feature");
+            NDList input = new NDList(array1, array2, array3, array4);
             // the NDList is mapped to (input1: list(Tensor), input2: Tensor)
             input.attach(manager);
 
-            IValue[] iValues = IValueUtils.getInputs(input);
+            Pair<IValue[], String> result = IValueUtils.getInputs(input);
+            IValue[] iValues = result.getKey();
             Assert.assertEquals(iValues.length, 2);
             Assert.assertTrue(iValues[0].isList());
             Assert.assertEquals(iValues[0].toIValueArray().length, 2);
+            Assert.assertEquals(result.getValue(), "get_text_feature");
 
             Arrays.stream(iValues).forEach(IValue::close);
         }