Repository: deepjavalibrary/djl
Pull request #1771 (2022-07-05 20:53:12)

[ort] Add support for registerCustomOpLibrary.
Change-Id: Ic3ec4249adb432b2a54668b7a45481698fa83ea5

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

```diff
@@             Coverage Diff              @@
##             master    #1771      +/-   ##
============================================
- Coverage     72.08%   70.67%   -1.42%     
- Complexity     5126     5575     +449     
============================================
  Files           473      527      +54     
  Lines         21970    24740    +2770     
  Branches       2351     2682     +331     
============================================
+ Hits          15838    17485    +1647     
- Misses         4925     5924     +999     
- Partials       1207     1331     +124     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1771?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/1771/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/1771/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/1771/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/1771/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/1771/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/1771/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/1771/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/1771/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/1771/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/1771/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 [411 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1771/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/1771?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/1771?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [0eaf1b5...370d351](https://codecov.io/gh/deepjavalibrary/djl/pull/1771?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/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
index 84c1f38f2..6f2cf648b 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtModel.java
@@ -179,6 +179,11 @@ public class OrtModel extends BaseModel {
             ortSession.setCPUArenaAllocator(true);
         }
 
+        String customOpLibrary = (String) options.get("customOpLibrary");
+        if (customOpLibrary != null) {
+            ortSession.registerCustomOpLibrary(customOpLibrary);
+        }
+
         Device device = manager.getDevice();
         if (options.containsKey("ortDevice")) {
             String ortDevice = (String) options.get("ortDevice");
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/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 36153510c..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);
@@ -607,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_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()


Test patch:
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);
         }
     }