Repository: deepjavalibrary/djl
Pull request #1719 (2022-06-16 05:19:28)

Advanced indexing that supports all indexing features on PyTorch
## Description ##

This PR enables all indexing features in pytorch. The indexing behaviour is guaranteed to be consistent, especially when mixed indices are passed.

The pytorch doc:
[C indexing API](https://pytorch.org/cppdocs/notes/tensor_indexing.html)
[source](https://pytorch.org/cppdocs/api/classat_1_1_tensor.html#_CPPv4NK2at6Tensor5indexE8ArrayRefIN2at8indexing11TensorIndexEE)

The relevant preceeding PRs: 
[Support of take from pytorch #1627](https://github.com/deepjavalibrary/djl/pull/1627)
[Add support of take on MXNet engine #1649
](https://github.com/deepjavalibrary/djl/pull/1649)


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

```diff
@@             Coverage Diff              @@
##             master    #1719      +/-   ##
============================================
- Coverage     72.08%   70.51%   -1.58%     
- Complexity     5126     5554     +428     
============================================
  Files           473      527      +54     
  Lines         21970    24690    +2720     
  Branches       2351     2679     +328     
============================================
+ Hits          15838    17411    +1573     
- Misses         4925     5948    +1023     
- Partials       1207     1331     +124     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1719?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/1719/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/1719/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/1719/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/1719/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/1719/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/1719/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/1719/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/1719/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/1719/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/1719/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 [413 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1719/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/1719?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/1719?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [22eb0d9...072c39a](https://codecov.io/gh/deepjavalibrary/djl/pull/1719?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/build.gradle b/api/build.gradle
index 5f16a668e..6e29aa793 100644
--- a/api/build.gradle
+++ b/api/build.gradle
@@ -9,6 +9,7 @@ dependencies {
     }
     testImplementation "org.slf4j:slf4j-simple:${slf4j_version}"
     testRuntimeOnly project(":engines:pytorch:pytorch-model-zoo")
+    testRuntimeOnly project(":engines:pytorch:pytorch-jni")
 }
 
 javadoc {
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index da5439be7..77162907f 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -530,11 +530,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @return the partial {@code NDArray}
      */
     default NDArray get(NDArray index) {
-        if (index.getDataType() == DataType.BOOLEAN) {
-            return get(new NDIndex().addBooleanIndex(index));
-        } else {
-            return take(index);
-        }
+        return get(new NDIndex("{}", index));
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index cdf76a659..f9bac2fb1 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -165,7 +165,7 @@ public interface NDManager extends AutoCloseable {
     ByteBuffer allocateDirect(int capacity);
 
     /**
-     * Creates a new {@code NDArray} if the input {@link NDArray} is from external engine.
+     * Creates a new {@code NDArray} if the input {@link NDArray} is from an external engine.
      *
      * @param array the input {@code NDArray}
      * @return a new {@code NDArray} if the input {@code NDArray} is from external engine
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 5eef36e45..8b8b40146 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDIndex.java
@@ -17,8 +17,10 @@ import ai.djl.ndarray.index.dim.NDIndexAll;
 import ai.djl.ndarray.index.dim.NDIndexBooleans;
 import ai.djl.ndarray.index.dim.NDIndexElement;
 import ai.djl.ndarray.index.dim.NDIndexFixed;
+import ai.djl.ndarray.index.dim.NDIndexNull;
 import ai.djl.ndarray.index.dim.NDIndexPick;
 import ai.djl.ndarray.index.dim.NDIndexSlice;
+import ai.djl.ndarray.index.dim.NDIndexTake;
 import ai.djl.ndarray.types.DataType;
 
 import java.util.ArrayList;
@@ -50,7 +52,7 @@ public class NDIndex {
     /* Android regex requires escape } char as well */
     private static final Pattern ITEM_PATTERN =
             Pattern.compile(
-                    "(\\*)|((-?\\d+|\\{\\})?:(-?\\d+|\\{\\})?(:(-?\\d+|\\{\\}))?)|(-?\\d+|\\{\\})");
+                    "(\\*)|((-?\\d+|\\{\\})?:(-?\\d+|\\{\\})?(:(-?\\d+|\\{\\}))?)|(-?\\d+|\\{\\})|null");
 
     private int rank;
     private List<NDIndexElement> indices;
@@ -105,6 +107,10 @@ public class NDIndex {
      *
      *     // Uses ellipsis to select all the dimensions except for last axis where we only get a subsection.
      *     assertEquals(a.get(new NDIndex("..., 2")).getShape(), new Shape(5, 4));
+     *
+     *     // 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));
+     *
      * </pre>
      *
      * @param indices a comma separated list of indices corresponding to either subsections,
@@ -335,6 +341,11 @@ public class NDIndex {
         if (!m.matches()) {
             throw new IllegalArgumentException("Invalid argument index: " + indexItem);
         }
+        // "null" case
+        if ("null".equals(indexItem)) {
+            indices.add(new NDIndexNull());
+            return argIndex;
+        }
         // "*" case
         String star = m.group(1);
         if (star != null) {
@@ -358,9 +369,12 @@ public class NDIndex {
                         indices.add(new NDIndexBooleans(array));
                         return argIndex + 1;
                     } else if (array.getDataType().isInteger()) {
-                        indices.add(new NDIndexPick(array));
+                        indices.add(new NDIndexTake(array));
                         return argIndex + 1;
                     }
+                } else if (arg == null) {
+                    indices.add(new NDIndexNull());
+                    return argIndex + 1;
                 }
                 throw new IllegalArgumentException("Unknown argument: " + arg);
             } else {
diff --git a/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexNull.java b/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexNull.java
new file mode 100644
index 000000000..d56182475
--- /dev/null
+++ b/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexNull.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.ndarray.index.dim;
+
+/** An {@code NDIndexElement} to return all values in a particular dimension. */
+public class NDIndexNull implements NDIndexElement {
+
+    /** {@inheritDoc} */
+    @Override
+    public int getRank() {
+        return 1;
+    }
+}
diff --git a/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexPick.java b/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexPick.java
index 4786353dc..ccd13b20e 100644
--- a/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexPick.java
+++ b/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexPick.java
@@ -17,15 +17,15 @@ import ai.djl.ndarray.NDArray;
 /** An {@link NDIndexElement} that gets elements by index in the specified axis. */
 public class NDIndexPick implements NDIndexElement {
 
-    private NDArray indices;
+    private NDArray index;
 
     /**
      * Constructs a pick.
      *
-     * @param indices the indices to pick
+     * @param index the index to pick
      */
-    public NDIndexPick(NDArray indices) {
-        this.indices = indices;
+    public NDIndexPick(NDArray index) {
+        this.index = index;
     }
 
     /** {@inheritDoc} */
@@ -35,11 +35,11 @@ public class NDIndexPick implements NDIndexElement {
     }
 
     /**
-     * Returns the indices to pick.
+     * Returns the index to pick.
      *
-     * @return the indices to pick
+     * @return the index to pick
      */
-    public NDArray getIndices() {
-        return indices;
+    public NDArray getIndex() {
+        return index;
     }
 }
diff --git a/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexTake.java b/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexTake.java
new file mode 100644
index 000000000..5ce76fe3a
--- /dev/null
+++ b/api/src/main/java/ai/djl/ndarray/index/dim/NDIndexTake.java
@@ -0,0 +1,45 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.ndarray.index.dim;
+
+import ai.djl.ndarray.NDArray;
+
+/** An {@link NDIndexElement} that gets elements by index in the specified axis. */
+public class NDIndexTake implements NDIndexElement {
+
+    private NDArray index;
+
+    /**
+     * Constructs a pick.
+     *
+     * @param index the index to pick
+     */
+    public NDIndexTake(NDArray index) {
+        this.index = index;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public int getRank() {
+        return 1;
+    }
+
+    /**
+     * Returns the index to pick.
+     *
+     * @return the index to pick
+     */
+    public NDArray getIndex() {
+        return index;
+    }
+}
diff --git a/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java b/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java
index 0136e895d..30d67bbba 100644
--- a/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java
+++ b/api/src/main/java/ai/djl/ndarray/index/full/NDIndexFullPick.java
@@ -53,7 +53,7 @@ public final class NDIndexFullPick {
                 axis++;
             } else if (el instanceof NDIndexPick) {
                 if (fullPick == null) {
-                    fullPick = new NDIndexFullPick(((NDIndexPick) el).getIndices(), axis);
+                    fullPick = new NDIndexFullPick(((NDIndexPick) el).getIndex(), axis);
                 } else {
                     // Don't support multiple picks
                     throw new UnsupportedOperationException(
diff --git a/engines/dlr/dlr-engine/build.gradle b/engines/dlr/dlr-engine/build.gradle
index 42401f9a3..45fe1852c 100644
--- a/engines/dlr/dlr-engine/build.gradle
+++ b/engines/dlr/dlr-engine/build.gradle
@@ -9,6 +9,7 @@ dependencies {
     }
     testImplementation "org.slf4j:slf4j-simple:${slf4j_version}"
     testRuntimeOnly project(":engines:pytorch:pytorch-engine")
+    testRuntimeOnly project(":engines:pytorch:pytorch-jni")
 }
 
 compileJava.dependsOn(processResources)
diff --git a/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle b/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle
index c23f27184..a04fa3640 100644
--- a/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle
+++ b/engines/paddlepaddle/paddlepaddle-model-zoo/build.gradle
@@ -10,6 +10,7 @@ dependencies {
     testImplementation "org.slf4j:slf4j-simple:${slf4j_version}"
     testImplementation(project(":testing"))
     testRuntimeOnly project(":engines:pytorch:pytorch-engine")
+    testRuntimeOnly project(":engines:pytorch:pytorch-jni")
 }
 
 task syncS3(type: Exec) {
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 9ffd96c12..07daec093 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
@@ -14,6 +14,7 @@ package ai.djl.pytorch.engine;
 
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.index.NDArrayIndexer;
+import ai.djl.ndarray.index.NDIndex;
 import ai.djl.ndarray.index.dim.NDIndexBooleans;
 import ai.djl.ndarray.index.full.NDIndexFullPick;
 import ai.djl.ndarray.index.full.NDIndexFullSlice;
@@ -49,6 +50,28 @@ public class PtNDArrayIndexer extends NDArrayIndexer {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray get(NDArray array, NDIndex index) {
+        if (index.getRank() == 0) {
+            if (array.getShape().isScalar()) {
+                return array.getManager() == manager
+                        ? array.duplicate()
+                        : manager.create(
+                                array.toByteBuffer(), array.getShape(), array.getDataType());
+            }
+            index.addAllDim();
+        }
+
+        if (array == null || array instanceof PtNDArray && array.getManager() == manager) {
+            return JniUtils.indexAdv((PtNDArray) array, index);
+        } else {
+            PtNDArray arrayNew =
+                    manager.create(array.toByteBuffer(), array.getShape(), array.getDataType());
+            return JniUtils.indexAdv(arrayNew, index);
+        }
+    }
+
     /** {@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 ebd5a998d..8e97adde6 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
@@ -14,6 +14,16 @@ package ai.djl.pytorch.jni;
 
 import ai.djl.Device;
 import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.index.dim.NDIndexAll;
+import ai.djl.ndarray.index.dim.NDIndexBooleans;
+import ai.djl.ndarray.index.dim.NDIndexElement;
+import ai.djl.ndarray.index.dim.NDIndexFixed;
+import ai.djl.ndarray.index.dim.NDIndexNull;
+import ai.djl.ndarray.index.dim.NDIndexPick;
+import ai.djl.ndarray.index.dim.NDIndexSlice;
+import ai.djl.ndarray.index.dim.NDIndexTake;
+import ai.djl.ndarray.index.full.NDIndexFullPick;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.ndarray.types.SparseFormat;
@@ -35,6 +45,8 @@ import java.nio.ByteOrder;
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.List;
+import java.util.ListIterator;
 import java.util.Set;
 
 /**
@@ -337,6 +349,69 @@ public final class JniUtils {
                         ndArray.getHandle(), minIndices, maxIndices, stepIndices));
     }
 
+    @SuppressWarnings("OptionalGetWithoutIsPresent")
+    public static PtNDArray indexAdv(PtNDArray ndArray, NDIndex index) {
+        if (ndArray == null) {
+            return ndArray;
+        }
+        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();
+                return pick(
+                        ndArray,
+                        ndArray.getManager().from(fullPick.getIndices()),
+                        fullPick.getAxis());
+            }
+        }
+        if (indices.size() == index.getEllipsisIndex()) {
+            PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+        }
+
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchIndexReturn(ndArray.getHandle(), torchIndexHandle));
+    }
+
     public static void indexSet(
             PtNDArray ndArray,
             PtNDArray value,
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 7bee9f7ab..11d2a95a5 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
@@ -600,4 +600,17 @@ final class PyTorchLibrary {
     native long torchNorm(long handle, int ord, long[] axis, boolean keepDims);
 
     native long torchNonZeros(long handle);
+
+    native long torchIndexInit(int size);
+
+    native long torchIndexReturn(long handle, long torchIndexHandle);
+
+    native void torchIndexAppendNoneEllipsis(long torchIndexHandle, boolean isEllipsis);
+
+    native void torchIndexAppendSlice(
+            long torchIndexHandle, long min, long max, long step, int nullSliceBinary);
+
+    native void torchIndexAppendFixed(long torchIndexHandle, long idx);
+
+    native void torchIndexAppendArray(long torchIndexHandle, long arrayHandle);
 }
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 ea48f69f1..484940106 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
@@ -121,6 +121,69 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndex(JNIEnv
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexInit(JNIEnv* env, jobject jthis, jint jsize) {
+  API_BEGIN()
+  std::vector<torch::indexing::TensorIndex>* index_ptr = new std::vector<torch::indexing::TensorIndex>;
+  index_ptr->reserve(jsize);
+  return reinterpret_cast<uintptr_t>(index_ptr);
+  API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexReturn(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlong jtorch_index_handle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  torch::Tensor* ret_ptr = new torch::Tensor(tensor_ptr->index(*index_ptr));
+  return reinterpret_cast<uintptr_t>(ret_ptr);
+  API_END_RETURN()
+}
+
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAppendNoneEllipsis(
+    JNIEnv* env, jobject jthis, jlong jtorch_index_handle, jboolean jis_ellipsis) {
+  API_BEGIN()
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  if (jis_ellipsis) {
+    index_ptr->emplace_back(torch::indexing::Ellipsis);
+  } else {
+    index_ptr->emplace_back(torch::indexing::None);
+  }
+  API_END()
+}
+
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAppendSlice(JNIEnv* env, jobject jthis,
+    jlong jtorch_index_handle, jlong jmin, jlong jmax, jlong jstep, jint jnull_slice_binary) {
+  API_BEGIN()
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  if (jnull_slice_binary == 0) {
+    index_ptr->emplace_back(torch::indexing::Slice(jmin, jmax, jstep));
+  } else if (jnull_slice_binary == 1) {
+    index_ptr->emplace_back(torch::indexing::Slice(jmin, torch::indexing::None, jstep));
+  } else if (jnull_slice_binary == 2) {
+    index_ptr->emplace_back(torch::indexing::Slice(torch::indexing::None, jmax, jstep));
+  } else if (jnull_slice_binary == 3) {
+    index_ptr->emplace_back(torch::indexing::Slice(torch::indexing::None, torch::indexing::None, jstep));
+  }
+  API_END()
+}
+
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAppendFixed(
+    JNIEnv* env, jobject jthis, jlong jtorch_index_handle, jlong jidx) {
+  API_BEGIN()
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  index_ptr->emplace_back((int) jidx);
+  API_END()
+}
+
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexAppendArray(
+    JNIEnv* env, jobject jthis, jlong jtorch_index_handle, jlong jarray) {
+  API_BEGIN()
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  auto* array_ptr = reinterpret_cast<torch::Tensor*>(jarray);
+  index_ptr->emplace_back(*array_ptr);
+  API_END()
+}
+
 JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIndexPut(JNIEnv* env, jobject jthis, jlong jhandle,
     jlong jvalue_handle, jlongArray jmin_indices, jlongArray jmax_indices, jlongArray jstep_indices) {
   API_BEGIN()
diff --git a/extensions/audio/build.gradle b/extensions/audio/build.gradle
index 8b6d60530..266cb9dd4 100644
--- a/extensions/audio/build.gradle
+++ b/extensions/audio/build.gradle
@@ -31,6 +31,7 @@ dependencies {
     testImplementation project(":testing")
     testRuntimeOnly "org.apache.logging.log4j:log4j-slf4j-impl:${log4j_slf4j_version}"
     testRuntimeOnly project(":engines:pytorch:pytorch-engine")
+    testRuntimeOnly project(":engines:pytorch:pytorch-jni")
 }
 
 publishing {
diff --git a/extensions/opencv/build.gradle b/extensions/opencv/build.gradle
index 7288cccbd..05dddad24 100644
--- a/extensions/opencv/build.gradle
+++ b/extensions/opencv/build.gradle
@@ -11,6 +11,7 @@ dependencies {
     testRuntimeOnly "org.apache.logging.log4j:log4j-slf4j-impl:${log4j_slf4j_version}"
 
     testRuntimeOnly project(":engines:pytorch:pytorch-model-zoo")
+    testRuntimeOnly project(":engines:pytorch:pytorch-jni")
 }
 
 publishing {
diff --git a/extensions/tablesaw/build.gradle b/extensions/tablesaw/build.gradle
index c3ed876da..2e71879f9 100644
--- a/extensions/tablesaw/build.gradle
+++ b/extensions/tablesaw/build.gradle
@@ -15,6 +15,7 @@ dependencies {
 //    testRuntimeOnly "tech.tablesaw:tablesaw-html:${tablesaw_version}"
 //    testRuntimeOnly "tech.tablesaw:tablesaw-json:${tablesaw_version}"
     testRuntimeOnly project(":engines:pytorch:pytorch-engine")
+    testRuntimeOnly project(":engines:pytorch:pytorch-jni")
 }
 
 publishing {


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 6e4368760..7d9f79fc8 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
@@ -120,10 +120,27 @@ public class NDIndexTest {
             Assert.assertEquals(original.get(bool), expected);
 
             // get from int array
-            original = manager.arange(1, 7f).reshape(-1, 3);
-            NDArray index = manager.create(new long[] {0, 4, 1, 2}, new Shape(2, 2));
+            original = manager.arange(1, 7f).reshape(-1, 2);
+            NDArray index = manager.create(new long[] {0, 0, 1, 2}, new Shape(2, 2));
             NDArray actual = original.get(index);
-            expected = manager.create(new float[] {1, 5, 2, 3}, new Shape(2, 2));
+            expected = manager.create(new float[] {1, 2, 1, 2, 3, 4, 5, 6}, new Shape(2, 2, 2));
+            Assert.assertEquals(actual, expected);
+
+            // indexing with boolean, broadcast int array and slice
+            original = manager.arange(3 * 3 * 3 * 3).reshape(3, 3, 3, 3);
+            NDArray bool1 = manager.create(new boolean[] {true, false, true});
+            NDArray index1 = manager.create(new long[] {2, 2}, new Shape(1, 2));
+            NDArray index2 = manager.create(new long[] {0, 1}, new Shape(1, 2));
+            actual = original.get(":{}, {}, {}, {}", 2, index1, bool1, index2);
+            expected = manager.create(new int[] {18, 25, 45, 52}, new Shape(2, 1, 2));
+            Assert.assertEquals(actual, expected);
+
+            // indexing with null, broadcast int array and slice
+            original = manager.arange(3 * 3 * 3).reshape(3, 3, 3);
+            index1 = manager.create(new long[] {0, 1}, new Shape(2));
+            index2 = manager.create(new long[] {0, 0, 2, 1}, new Shape(2, 2));
+            actual = original.get(":{}, {}, {}, {}", 2, index1, index2, null);
+            expected = manager.create(new int[] {0, 3, 2, 4, 9, 12, 11, 13}, new Shape(2, 2, 2, 1));
             Assert.assertEquals(actual, expected);
         }
     }