Repository: deepjavalibrary/djl
Pull request #1762 (2022-06-30 18:22:20)

[pytorch] Add NDList to IValue unit test
The testing code can be used as example as well

Change-Id: I693624fd111afb0a5547d619a03d2cb15831af1b

## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1762?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1762](https://codecov.io/gh/deepjavalibrary/djl/pull/1762?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (5f7d1e3) into [master](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) (bb5073f) will **decrease** coverage by `1.45%`.
> The diff coverage is `70.46%`.

```diff
@@             Coverage Diff              @@
##             master    #1762      +/-   ##
============================================
- Coverage     72.08%   70.62%   -1.46%     
- Complexity     5126     5567     +441     
============================================
  Files           473      527      +54     
  Lines         21970    24702    +2732     
  Branches       2351     2680     +329     
============================================
+ Hits          15838    17447    +1609     
- Misses         4925     5927    +1002     
- Partials       1207     1328     +121     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1762?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/1762/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/1762/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/1762/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/1762/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/1762/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/1762/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% <ø> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1762/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1762/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../cv/translator/StyleTransferTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1762/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU3R5bGVUcmFuc2ZlclRyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `40.00% <ø> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1762/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `8.33% <0.00%> (-0.50%)` | :arrow_down: |
| ... and [410 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1762/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1762?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1762?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [3655b89...5f7d1e3](https://codecov.io/gh/deepjavalibrary/djl/pull/1762?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 77162907f..c922d7e75 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -581,14 +581,24 @@ public interface NDArray extends NDResource, BytesSupplier {
     NDArray gather(NDArray index, int axis);
 
     /**
-     * Returns a partial {@code NDArray} pointed by the indexed array according to linear indexing,
-     * and the of output is of the same shape as index.
+     * Returns a partial {@code NDArray} pointed by index according to linear indexing, and the of
+     * output is of the same shape as index.
      *
-     * @param index picks the elements of an NDArray to the same position as index
+     * @param index picks the elements of an NDArray and output to the same entry as in index
      * @return the partial {@code NDArray} of the same shape as index
      */
     NDArray take(NDArray index);
 
+    /**
+     * Set the entries of {@code NDArray} pointed by index according to linear indexing, to be the
+     * numbers in data, which is of the same shape as index.
+     *
+     * @param index select the entries of an {@code NDArray}
+     * @param data numbers to assign to the indexed entries
+     * @return The NDArray with updated values
+     */
+    NDArray put(NDArray index, NDArray data);
+
     /**
      * Returns a scalar {@code NDArray} corresponding to a single element.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 0770bbeb5..ab619bf9c 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -189,6 +189,12 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException(UNSUPPORTED_MSG);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray put(NDArray index, NDArray data) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void set(Buffer data) {
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index d3936365d..1db6a819c 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -330,6 +330,12 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         return manager.invoke("take", new NDList(this.flatten(), index), params).singletonOrThrow();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray put(NDArray index, NDArray data) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void copyTo(NDArray ndArray) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index f0e04d35b..70654e6c4 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -258,6 +258,15 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.take(this, (PtNDArray) index);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray put(NDArray index, NDArray data) {
+        if (!(index instanceof PtNDArray) || !(data instanceof PtNDArray)) {
+            throw new IllegalArgumentException("Only PtNDArray is supported.");
+        }
+        return JniUtils.put(this, (PtNDArray) index, (PtNDArray) data);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void copyTo(NDArray array) {
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 8e97adde6..a4f2a6100 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
@@ -445,6 +445,16 @@ public final class JniUtils {
                 PyTorchLibrary.LIB.torchTake(ndArray.getHandle(), index.getHandle()));
     }
 
+    public static PtNDArray put(PtNDArray ndArray, PtNDArray index, PtNDArray data) {
+        if (index.getDataType() != DataType.INT64) {
+            index = index.toType(DataType.INT64, true);
+        }
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchPut(
+                        ndArray.getHandle(), index.getHandle(), data.getHandle()));
+    }
+
     public static PtNDArray pick(PtNDArray ndArray, PtNDArray index, long dim) {
         Shape indexShape = index.getShape();
         Shape ndShape = ndArray.getShape();
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 11d2a95a5..35593fb48 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
@@ -200,6 +200,8 @@ final class PyTorchLibrary {
 
     native long torchTake(long handle, long index);
 
+    native long torchPut(long handle, long index, long data);
+
     native long torchMaskedSelect(long handle, long maskHandle);
 
     native void torchMaskedPut(long handle, long valueHandle, long maskHandle);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
index 484940106..21be55ae5 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
@@ -249,6 +249,17 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchTake(
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchPut(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlong jindex_handle, jlong jdata_handle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* index_ptr = reinterpret_cast<torch::Tensor*>(jindex_handle);
+  const auto* data_ptr = reinterpret_cast<torch::Tensor*>(jdata_handle);
+  const auto* result_ptr = new torch::Tensor(tensor_ptr->put(*index_ptr, *data_ptr));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchMaskedSelect(
     JNIEnv* env, jobject jthis, jlong jhandle, jlong jmasked_handle) {
   API_BEGIN()
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index e1ff497bc..3a6973598 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -220,6 +220,12 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray put(NDArray index, NDArray data) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void attach(NDManager manager) {
diff --git a/extensions/fasttext/src/main/native/ai_djl_fasttext_jni_FastTextLibrary.cc b/extensions/fasttext/src/main/native/ai_djl_fasttext_jni_FastTextLibrary.cc
index 8a14c1078..faad19a86 100644
--- a/extensions/fasttext/src/main/native/ai_djl_fasttext_jni_FastTextLibrary.cc
+++ b/extensions/fasttext/src/main/native/ai_djl_fasttext_jni_FastTextLibrary.cc
@@ -116,14 +116,15 @@ JNIEXPORT jint JNICALL Java_ai_djl_fasttext_jni_FastTextLibrary_predictProba(
   fasttext_ptr->predictLine(in, predictions, top_k, 0.0);
 
   int size = predictions.size();
-   jclass java_lang_Float = static_cast<jclass>(env->NewGlobalRef(env->FindClass("java/lang/Float")));
-   jmethodID java_lang_Float_ = env->GetMethodID(java_lang_Float, "<init>", "(F)V");
+  jclass java_lang_Float = static_cast<jclass>(env->NewGlobalRef(env->FindClass("java/lang/Float")));
+  jmethodID java_lang_Float_ = env->GetMethodID(java_lang_Float, "<init>", "(F)V");
   jclass java_util_ArrayList = static_cast<jclass>(env->NewGlobalRef(env->FindClass("java/util/ArrayList")));
   jmethodID java_util_ArrayList_add = env->GetMethodID(java_util_ArrayList, "add", "(Ljava/lang/Object;)Z");
   for (int i = 0; i < size; ++i) {
     std::pair<real, std::string> pair = predictions[i];
-     env->CallBooleanMethod(jclasses, java_util_ArrayList_add, env->NewStringUTF(pair.second.c_str()));
-     env->CallBooleanMethod(jprob, java_util_ArrayList_add, env->NewObject(java_lang_Float, java_lang_Float_, pair.first));
+    env->CallBooleanMethod(jclasses, java_util_ArrayList_add, env->NewStringUTF(pair.second.c_str()));
+    env->CallBooleanMethod(
+        jprob, java_util_ArrayList_add, env->NewObject(java_lang_Float, java_lang_Float_, pair.first));
   }
 
   return size;
diff --git a/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc b/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc
index de0d94261..dd2c80ed5 100644
--- a/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc
+++ b/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc
@@ -38,13 +38,13 @@ JNIEXPORT void JNICALL Java_ai_djl_sentencepiece_jni_SentencePieceLibrary_loadMo
 }
 
 JNIEXPORT void JNICALL Java_ai_djl_sentencepiece_jni_SentencePieceLibrary_loadModelFromBytes(
-        JNIEnv* env, jobject jthis, jlong jhandle, jbyteArray jserialized) {
-    auto* processor_ptr = reinterpret_cast<sentencepiece::SentencePieceProcessor*>(jhandle);
-    int length = env->GetArrayLength(jserialized);
-    std::vector<char> buff(length, 0);
-    env->GetByteArrayRegion(jserialized, 0, length, reinterpret_cast<jbyte*>(buff.data()));
-    std::string serialized(buff.data(), buff.size());
-    CheckStatus(env, processor_ptr->LoadFromSerializedProto(serialized));
+    JNIEnv* env, jobject jthis, jlong jhandle, jbyteArray jserialized) {
+  auto* processor_ptr = reinterpret_cast<sentencepiece::SentencePieceProcessor*>(jhandle);
+  int length = env->GetArrayLength(jserialized);
+  std::vector<char> buff(length, 0);
+  env->GetByteArrayRegion(jserialized, 0, length, reinterpret_cast<jbyte*>(buff.data()));
+  std::string serialized(buff.data(), buff.size());
+  CheckStatus(env, processor_ptr->LoadFromSerializedProto(serialized));
 }
 
 JNIEXPORT void JNICALL Java_ai_djl_sentencepiece_jni_SentencePieceLibrary_deleteSentencePieceProcessor(


Test patch:
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 da3bdd15a..2afbb10c2 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
@@ -12,11 +12,10 @@
  */
 package ai.djl.pytorch.jni;
 
+import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
-import ai.djl.pytorch.engine.PtNDArray;
-import ai.djl.pytorch.engine.PtNDManager;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -24,14 +23,16 @@ import org.testng.annotations.Test;
 import java.util.Arrays;
 
 public class IValueUtilsTest {
+
     @Test
-    public void getInputsTestTupleSyntax() {
-        try (PtNDManager manager = (PtNDManager) NDManager.newBaseManager()) {
-            PtNDArray array1 = (PtNDArray) manager.zeros(new Shape(1));
-            array1.setName("Test()");
-            PtNDArray array2 = (PtNDArray) manager.ones(new Shape(1));
-            array2.setName("Test()");
+    public void testTuple() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array1 = manager.zeros(new Shape(1));
+            array1.setName("input1()");
+            NDArray array2 = manager.ones(new Shape(1));
+            array2.setName("input1()");
             NDList input = new NDList(array1, array2);
+            // the NDList is mapped to (input1: Tuple(Tensor))
             input.attach(manager);
 
             IValue[] iValues = IValueUtils.getInputs(input);
@@ -42,4 +43,54 @@ public class IValueUtilsTest {
             Arrays.stream(iValues).forEach(IValue::close);
         }
     }
+
+    @Test
+    public void testMapOfTensor() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array1 = manager.zeros(new Shape(1));
+            array1.setName("input1.key1");
+            NDArray array2 = manager.ones(new Shape(1));
+            array2.setName("input1.key2");
+            NDArray array3 = manager.zeros(new Shape(1));
+            array3.setName("input2.key1");
+            NDArray array4 = manager.ones(new Shape(1));
+            array4.setName("input2.key2");
+            NDArray array5 = manager.ones(new Shape(1));
+            array5.setName("input2.key3");
+            NDList input = new NDList(array1, array2, array3, array4, array5);
+            // the NDList is mapped to (input1: Dict(str, Tensor), input2: Dict(str, Tensor))
+            // the first part of NDArray name is the variable name of the inputs
+            // the 2nd part is the key of each value in the dict
+            input.attach(manager);
+
+            IValue[] iValues = IValueUtils.getInputs(input);
+            Assert.assertEquals(iValues.length, 2);
+            Assert.assertTrue(iValues[0].isMap());
+            Assert.assertEquals(iValues[1].toIValueMap().size(), 3);
+
+            Arrays.stream(iValues).forEach(IValue::close);
+        }
+    }
+
+    @Test
+    public void testListOfTensor() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array1 = manager.zeros(new Shape(1));
+            array1.setName("input1[]");
+            NDArray array2 = manager.ones(new Shape(1));
+            array2.setName("input1[]");
+            NDArray array3 = manager.ones(new Shape(1));
+            array3.setName("input2");
+            NDList input = new NDList(array1, array2, array3);
+            // the NDList is mapped to (input1: list(Tensor), input2: Tensor)
+            input.attach(manager);
+
+            IValue[] iValues = IValueUtils.getInputs(input);
+            Assert.assertEquals(iValues.length, 2);
+            Assert.assertTrue(iValues[0].isList());
+            Assert.assertEquals(iValues[0].toIValueArray().length, 2);
+
+            Arrays.stream(iValues).forEach(IValue::close);
+        }
+    }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
index 7d9f79fc8..49c165fd1 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
@@ -228,4 +228,15 @@ public class NDIndexTest {
             Assert.assertEquals(original, expected);
         }
     }
+
+    @Test
+    public void testPut() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray original = manager.create(new float[] {1, 2, 3, 4}, new Shape(2, 2));
+            NDArray expected = manager.create(new float[] {1, 8, 666, 77}, new Shape(2, 2));
+            NDArray idx = manager.create(new long[] {2, 3, 1}, new Shape(3));
+            NDArray data = manager.create(new float[] {666, 77, 8}, new Shape(3));
+            Assert.assertEquals(original.put(idx, data), expected);
+        }
+    }
 }