Repository: deepjavalibrary/djl
Pull request #1755 (2022-06-28 18:05:42)

Add tensor setter feature with advanced indexing on PyTorch engine
## Description ##

Add tensor setter feature with advanced indexing on PyTorch engine. This is the counterpart of  the getter feature in one previous PR. [Advanced indexing that supports all indexing features on PyTorch #1719](https://github.com/deepjavalibrary/djl/pull/1719)

This PR succeeds [Add put feature with linear indexing on PyTorch engine #1749](https://github.com/deepjavalibrary/djl/pull/1749).


The relavant preceding PRs
[Advanced indexing that supports all indexing features on PyTorch #1719](https://github.com/deepjavalibrary/djl/pull/1719)
[Add put feature with linear indexing on PyTorch engine #1749](https://github.com/deepjavalibrary/djl/pull/1749)

codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1755?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1755](https://codecov.io/gh/deepjavalibrary/djl/pull/1755?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (4974abf) 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.61%`.
> The diff coverage is `70.36%`.

```diff
@@             Coverage Diff              @@
##             master    #1755      +/-   ##
============================================
- Coverage     72.08%   70.47%   -1.62%     
- Complexity     5126     5574     +448     
============================================
  Files           473      527      +54     
  Lines         21970    24759    +2789     
  Branches       2351     2698     +347     
============================================
+ Hits          15838    17448    +1610     
- Misses         4925     5973    +1048     
- Partials       1207     1338     +131     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1755?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/1755/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/1755/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/1755/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/1755/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/1755/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/1755/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/1755/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/1755/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/1755/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/1755/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/1755/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/1755?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/1755?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [51b38b6...4974abf](https://codecov.io/gh/deepjavalibrary/djl/pull/1755?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 c922d7e75..516e279a2 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -595,7 +595,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      *
      * @param index select the entries of an {@code NDArray}
      * @param data numbers to assign to the indexed entries
-     * @return The NDArray with updated values
+     * @return the NDArray with updated values
      */
     NDArray put(NDArray index, NDArray data);
 
diff --git a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
index 47994b22c..3c7f67d4f 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
@@ -79,89 +79,80 @@ public abstract class NDArrayIndexer {
     }
 
     /**
-     * Sets the values of the array at the fullSlice with an array.
-     *
-     * @param array the array to set
-     * @param fullSlice the fullSlice of the index to set in the array
-     * @param value the value to set with
-     */
-    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, NDArray value);
-
-    /**
-     * Sets the values of the array at the boolean locations with an array.
-     *
-     * @param array the array to set
-     * @param indices a boolean array where true indicates values to update
-     * @param value the value to set with when condition is true
-     */
-    public void set(NDArray array, NDIndexBooleans indices, NDArray value) {
-        array.intern(NDArrays.where(indices.getIndex(), value, array));
-    }
-
-    /**
-     * Sets the values of the array at the index locations with an array.
+     * Sets the entries of array at the indexed locations with the parameter value. The value can be
+     * only Number or NDArray.
      *
      * @param array the array to set
      * @param index the index to set at in the array
      * @param value the value to set with
      */
-    public void set(NDArray array, NDIndex index, NDArray value) {
+    public void set(NDArray array, NDIndex index, Object value) {
+        NDIndexFullSlice fullSlice =
+                NDIndexFullSlice.fromIndex(index, array.getShape()).orElse(null);
+        if (fullSlice != null) {
+            if (value instanceof Number) {
+                set(array, fullSlice, (Number) value);
+            } else if (value instanceof NDArray) {
+                set(array, fullSlice, (NDArray) value);
+            } else {
+                throw new IllegalArgumentException(
+                        "The type of value to assign cannot be other than NDArray and Number.");
+            }
+            return;
+        }
+
         List<NDIndexElement> indices = index.getIndices();
         if (!indices.isEmpty() && indices.get(0) instanceof NDIndexBooleans) {
             if (indices.size() != 1) {
                 throw new IllegalArgumentException(
-                        "get() currently didn't support more that one boolean NDArray");
+                        "set() currently doesn't support more than one boolean NDArray");
+            }
+            if (value instanceof Number) {
+                set(
+                        array,
+                        (NDIndexBooleans) indices.get(0),
+                        array.getManager().create((Number) value));
+            } else if (value instanceof NDArray) {
+                set(array, (NDIndexBooleans) indices.get(0), (NDArray) value);
+            } else {
+                throw new IllegalArgumentException(
+                        "The type of value to assign cannot be other than NDArray and Number.");
             }
-            set(array, (NDIndexBooleans) indices.get(0), value);
         }
 
-        NDIndexFullSlice fullSlice =
-                NDIndexFullSlice.fromIndex(index, array.getShape()).orElse(null);
-        if (fullSlice != null) {
-            set(array, fullSlice, value);
-            return;
-        }
         throw new UnsupportedOperationException(
                 "set() currently supports all, fixed, and slices indices");
     }
 
     /**
-     * Sets the values of the array at the fullSlice with a number.
+     * Sets the values of the array at the fullSlice with an array.
      *
      * @param array the array to set
      * @param fullSlice the fullSlice of the index to set in the array
      * @param value the value to set with
      */
-    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, Number value);
+    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, NDArray value);
 
     /**
-     * Sets the values of the array at the index locations with a number.
+     * Sets the values of the array at the boolean locations with an array.
      *
      * @param array the array to set
-     * @param index the index to set at in the array
-     * @param value the value to set with
+     * @param indices a boolean array where true indicates values to update
+     * @param value the value to set with when condition is true
      */
-    public void set(NDArray array, NDIndex index, Number value) {
-        NDIndexFullSlice fullSlice =
-                NDIndexFullSlice.fromIndex(index, array.getShape()).orElse(null);
-        if (fullSlice != null) {
-            set(array, fullSlice, value);
-            return;
-        }
-        // use booleanMask for NDIndexBooleans case
-        List<NDIndexElement> indices = index.getIndices();
-        if (!indices.isEmpty() && indices.get(0) instanceof NDIndexBooleans) {
-            if (indices.size() != 1) {
-                throw new IllegalArgumentException(
-                        "set() currently didn't support more that one boolean NDArray");
-            }
-            set(array, (NDIndexBooleans) indices.get(0), array.getManager().create(value));
-            return;
-        }
-        throw new UnsupportedOperationException(
-                "set() currently supports all, fixed, and slices indices");
+    public void set(NDArray array, NDIndexBooleans indices, NDArray value) {
+        array.intern(NDArrays.where(indices.getIndex(), value, array));
     }
 
+    /**
+     * Sets the values of the array at the fullSlice with a number.
+     *
+     * @param array the array to set
+     * @param fullSlice the fullSlice of the index to set in the array
+     * @param value the value to set with
+     */
+    public abstract void set(NDArray array, NDIndexFullSlice fullSlice, Number value);
+
     /**
      * Sets a scalar value in the array at the indexed location.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/index/NDIndex.java b/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
index 8b8b40146..a5fb0a2f1 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
@@ -111,6 +111,11 @@ public class NDIndex {
      *     // Uses null to add an extra axis to the output array
      *     assertEquals(a.get(new NDIndex(":2, null, 0, :2")).getShape(), new Shape(2, 1, 2));
      *
+     *     // Gets entries of an NDArray with mixed index
+     *     index1 = manager.create(new long[] {0, 1, 1}, new Shape(2));
+     *     bool1 = manager.create(new boolean[] {true, false, true});
+     *     assertEquals(a.get(new NDIndex(":{}, {}, {}, {}" 2, index1, bool1, null).getShape(), new Shape(2, 2, 1));
+     *
      * </pre>
      *
      * @param indices a comma separated list of indices corresponding to either subsections,
diff --git a/docs/development/setup.md b/docs/development/setup.md
index 709c4a8a8..168263e59 100644
--- a/docs/development/setup.md
+++ b/docs/development/setup.md
@@ -2,7 +2,7 @@
 
 ## Install the Java Development Kit
 
-The Deep Java Library (DJL)  project requires JDK 8 (or later). We recommend using JDK8, since there are some known issues with JDK 11+.
+The Deep Java Library (DJL)  project requires JDK 11 (or later). 
 
 Verify that Java is available in your $PATH environment variable by using the following commands. If you have multiple versions of Java installed,
 you can use the $JAVA_HOME environment variable to control which version of Java to use.
@@ -10,13 +10,13 @@ you can use the $JAVA_HOME environment variable to control which version of Java
 For ubuntu:
 
 ```bash
-sudo apt-get install openjdk-8-jdk-headless
+sudo apt-get install openjdk-11-jdk
 ```
 
 For centos
 
 ```bash
-sudo yum install java-1.8.0-openjdk
+sudo yum install java-11-openjdk
 ```
 
 For Mac:
@@ -24,7 +24,7 @@ For Mac:
 ```bash
 brew tap homebrew/cask-versions
 brew update
-brew install --cask temurin8
+brew install --cask temurin11
 ```
 
 You can also download and install [Oracle JDK](https://www.oracle.com/technetwork/java/javase/overview/index.html)
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
index 07daec093..fcc8958d0 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
@@ -72,6 +72,25 @@ public class PtNDArrayIndexer extends NDArrayIndexer {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void set(NDArray array, NDIndex index, Object data) {
+        PtNDArray ptArray =
+                array instanceof PtNDArray
+                        ? (PtNDArray) array
+                        : manager.create(
+                                array.toByteBuffer(), array.getShape(), array.getDataType());
+
+        if (data instanceof Number) {
+            JniUtils.indexAdvPut(ptArray, index, (PtNDArray) manager.create((Number) data));
+        } else if (data instanceof NDArray) {
+            JniUtils.indexAdvPut(ptArray, index, (PtNDArray) data);
+        } else {
+            throw new IllegalArgumentException(
+                    "The type of value to assign cannot be other than NDArray and Number.");
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public void set(NDArray array, NDIndexFullSlice fullSlice, NDArray value) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
index 1cbd438d1..79cb3f577 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValue.java
@@ -287,6 +287,24 @@ public class IValue extends NativeResource<Long> {
         return new IValue(PyTorchLibrary.LIB.iValueFromStringMap(keys, handles));
     }
 
+    /**
+     * Creates a new {@code IValue} of type {@code Map[String, IValue]}.
+     *
+     * @param map the Map[String, IValue] value
+     * @return a new {@code IValue} of type {@code Map[String, IValue]}
+     */
+    public static IValue stringIValueMapFrom(Map<String, IValue> map) {
+        String[] keys = new String[map.size()];
+        long[] handles = new long[map.size()];
+        int i = 0;
+        for (Map.Entry<String, IValue> entry : map.entrySet()) {
+            keys[i] = entry.getKey();
+            handles[i] = entry.getValue().getHandle();
+            ++i;
+        }
+        return new IValue(PyTorchLibrary.LIB.iValueFromStringIValueMap(keys, handles));
+    }
+
     /**
      * Returns the {@code boolean} value of this IValue.
      *
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 a4f2a6100..3364cd45e 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
@@ -409,7 +409,69 @@ public final class JniUtils {
 
         return new PtNDArray(
                 ndArray.getManager(),
-                PyTorchLibrary.LIB.torchIndexReturn(ndArray.getHandle(), torchIndexHandle));
+                PyTorchLibrary.LIB.torchIndexAdvGet(ndArray.getHandle(), torchIndexHandle));
+    }
+
+    @SuppressWarnings("OptionalGetWithoutIsPresent")
+    public static void indexAdvPut(PtNDArray ndArray, NDIndex index, PtNDArray data) {
+        if (ndArray == null) {
+            return;
+        }
+
+        // Index aggregation
+        List<NDIndexElement> indices = index.getIndices();
+        long torchIndexHandle = PyTorchLibrary.LIB.torchIndexInit(indices.size());
+        ListIterator<NDIndexElement> it = indices.listIterator();
+        while (it.hasNext()) {
+            if (it.nextIndex() == index.getEllipsisIndex()) {
+                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+            }
+
+            NDIndexElement elem = it.next();
+            if (elem instanceof NDIndexNull) {
+                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, false);
+            } else if (elem instanceof NDIndexSlice) {
+                Long min = ((NDIndexSlice) elem).getMin();
+                Long max = ((NDIndexSlice) elem).getMax();
+                Long step = ((NDIndexSlice) elem).getStep();
+                int nullSliceBin = (min == null ? 1 : 0) * 2 + (max == null ? 1 : 0);
+                // nullSliceBin encodes whether the slice (min, max) is null:
+                // is_null == 1, ! is_null == 0;
+                // 0b11 == 3, 0b10 = 2, ...
+                PyTorchLibrary.LIB.torchIndexAppendSlice(
+                        torchIndexHandle,
+                        min == null ? 0 : min,
+                        max == null ? 0 : max,
+                        step == null ? 1 : step,
+                        nullSliceBin);
+            } else if (elem instanceof NDIndexAll) {
+                PyTorchLibrary.LIB.torchIndexAppendSlice(torchIndexHandle, 0, 0, 1, 3);
+            } else if (elem instanceof NDIndexFixed) {
+                PyTorchLibrary.LIB.torchIndexAppendFixed(
+                        torchIndexHandle, ((NDIndexFixed) elem).getIndex());
+            } else if (elem instanceof NDIndexBooleans) {
+                PtNDArray indexArr = (PtNDArray) ((NDIndexBooleans) elem).getIndex();
+                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
+            } else if (elem instanceof NDIndexTake) {
+                PtNDArray indexArr = (PtNDArray) ((NDIndexTake) elem).getIndex();
+                if (indexArr.getDataType() != DataType.INT64) {
+                    indexArr = indexArr.toType(DataType.INT64, true);
+                }
+                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
+            } else if (elem instanceof NDIndexPick) {
+                // Backward compatible
+                NDIndexFullPick fullPick =
+                        NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
+                pick(ndArray, ndArray.getManager().from(fullPick.getIndices()), fullPick.getAxis());
+                return;
+            }
+        }
+        if (indices.size() == index.getEllipsisIndex()) {
+            PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+        }
+
+        PyTorchLibrary.LIB.torchIndexAdvPut(
+                ndArray.getHandle(), torchIndexHandle, data.getHandle());
     }
 
     public static void indexSet(
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 35593fb48..542ef3b54 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
@@ -192,6 +192,8 @@ final class PyTorchLibrary {
             long[] maxIndices,
             long[] stepIndices);
 
+    native void torchIndexAdvPut(long handle, long torchIndexHandle, long data);
+
     native void torchSet(long handle, ByteBuffer data);
 
     native long torchSlice(long handle, long dim, long start, long end, long step);
@@ -524,6 +526,8 @@ final class PyTorchLibrary {
 
     native long iValueFromStringMap(String[] keys, long[] tensorHandles);
 
+    native long iValueFromStringIValueMap(String[] keys, long[] tensorHandles);
+
     native long iValueToTensor(long iValueHandle);
 
     native boolean iValueToBool(long iValueHandle);
@@ -605,7 +609,7 @@ final class PyTorchLibrary {
 
     native long torchIndexInit(int size);
 
-    native long torchIndexReturn(long handle, long torchIndexHandle);
+    native long torchIndexAdvGet(long handle, long torchIndexHandle);
 
     native void torchIndexAppendNoneEllipsis(long torchIndexHandle, boolean isEllipsis);
 
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc
index 5a7b40171..920565f7c 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_ivalue.cc
@@ -176,6 +176,30 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueFromStringM
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueFromStringIValueMap(
+    JNIEnv* env, jobject jthis, jobjectArray jkeys, jlongArray jvalues) {
+  API_BEGIN()
+  auto len = static_cast<size_t>(env->GetArrayLength(jvalues));
+  jlong* jptrs = env->GetLongArrayElements(jvalues, JNI_FALSE);
+  if (len == 0) {
+    const auto* ivalue_ptr = new torch::IValue{c10::impl::GenericDict(c10::StringType::get(), c10::TensorType::get())};
+    return reinterpret_cast<uintptr_t>(ivalue_ptr);
+  }
+
+  auto* firstEntryValue = reinterpret_cast<torch::IValue*>(jptrs[0]);
+  c10::impl::GenericDict dict(c10::StringType::get(), c10::unshapedType(firstEntryValue->type()));
+  for (size_t i = 0; i < len; ++i) {
+    auto jname = (jstring) env->GetObjectArrayElement(jkeys, i);
+    std::string name = djl::utils::jni::GetStringFromJString(env, jname);
+    dict.insert(name, *reinterpret_cast<torch::IValue*>(jptrs[i]));
+  }
+  env->ReleaseLongArrayElements(jvalues, jptrs, JNI_ABORT);
+  env->DeleteLocalRef(jkeys);
+  const auto* ivalue_ptr = new torch::IValue{dict};
+  return reinterpret_cast<uintptr_t>(ivalue_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_iValueToTensor(
     JNIEnv* env, jobject jthis, jlong jhandle) {
   API_BEGIN()
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 21be55ae5..2b89fb98f 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
@@ -129,7 +129,7 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexInit(JN
   API_END_RETURN()
 }
 
-JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexReturn(
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAdvGet(
     JNIEnv* env, jobject jthis, jlong jhandle, jlong jtorch_index_handle) {
   API_BEGIN()
   const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
@@ -194,6 +194,16 @@ JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexPut(JNIE
   API_END()
 }
 
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAdvPut(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlong jtorch_index_handle, jlong jdata_handle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const auto* data_ptr = reinterpret_cast<torch::Tensor*>(jdata_handle);
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  ((torch::Tensor) *tensor_ptr).index_put_(*index_ptr, *data_ptr);
+  API_END()
+}
+
 JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchSet(
     JNIEnv* env, jobject jthis, jlong jhandle, jobject jbuffer) {
   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 3a6973598..a48f03aa5 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
@@ -116,13 +116,7 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
     /** {@inheritDoc} */
     @Override
     public SparseFormat getSparseFormat() {
-        throw new UnsupportedOperationException("Not implemented");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public boolean isSparse() {
-        throw new UnsupportedOperationException("Not implemented");
+        return SparseFormat.DENSE;
     }
 
     /** {@inheritDoc} */


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
index 56393f510..4caca064f 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/IValueTest.java
@@ -123,6 +123,18 @@ public class IValueTest {
                 Assert.assertEquals(list.get("data1"), array1);
             }
 
+            // (Dict(str, Tensor[])
+            Map<String, IValue> iValueMap = new ConcurrentHashMap<>();
+            try (IValue v1 = IValue.listFrom(array1);
+                    IValue v2 = IValue.listFrom(array2)) {
+                iValueMap.put("data1", v1);
+                iValueMap.put("data2", v2);
+                try (IValue ivalue = IValue.stringIValueMapFrom(iValueMap)) {
+                    Assert.assertTrue(ivalue.isMap());
+                    Assert.assertEquals(ivalue.getType(), "Dict(str, Tensor[])");
+                }
+            }
+
             try (IValue iv1 = IValue.from(1);
                     IValue iv2 = IValue.from(2);
                     IValue ivalue = IValue.listFrom(iv1, iv2)) {
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 49c165fd1..c9c6744bd 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
@@ -157,6 +157,15 @@ public class NDIndexTest {
             expected = manager.create(new int[] {0, 1, 9, 10, 4, 5, 11, 12}, new Shape(2, 4));
             original.set(new NDIndex(":, 2:"), manager.arange(9, 13).reshape(2, 2));
             Assert.assertEquals(original, expected);
+
+            // set by index array
+            original = manager.arange(1, 10).reshape(3, 3);
+            NDArray index = manager.create(new long[] {0, 1}, new Shape(2));
+            value = manager.create(new int[] {666, 777, 888, 999}, new Shape(2, 2));
+            original.set(new NDIndex("{}, :{}", index, 2), value);
+            expected =
+                    manager.create(new int[] {666, 777, 3, 888, 999, 6, 7, 8, 9}, new Shape(3, 3));
+            Assert.assertEquals(original, expected);
         }
     }
 
@@ -188,6 +197,14 @@ public class NDIndexTest {
             expected = manager.create(new float[] {1, 1, 1, 3}).reshape(2, 2);
             original.set(new NDIndex("..., 0"), 1);
             Assert.assertEquals(original, expected);
+
+            // set by index array
+            original = manager.arange(1, 10).reshape(3, 3);
+            NDArray index = manager.create(new long[] {0, 1}, new Shape(2));
+            original.set(new NDIndex("{}, :{}", index, 2), 666);
+            expected =
+                    manager.create(new int[] {666, 666, 3, 666, 666, 6, 7, 8, 9}, new Shape(3, 3));
+            Assert.assertEquals(original, expected);
         }
     }