Repository: deepjavalibrary/djl
Pull request #1622 (2022-05-03 23:02:12)

Support gather of pytorch
## Description ##

This originiates from issue https://github.com/deepjavalibrary/djl/issues/248. This is regarding support of advanced indexing.

As part of this support, this PR implements `gather`. It's the same as the gather in
https://pytorch.org/docs/1.11/generated/torch.gather.html#torch.gather 
https://numpy.org/doc/stable/reference/generated/numpy.take_along_axis.html
It currently supports pytorch engine only. MXNet and TensorFlow have different defination of gather.

The design is the following:
`gather(arr, idx, axis)` returns the following based on the conventions [1](https://pytorch.org/docs/1.11/generated/torch.gather.html#torch.gather) [2](https://numpy.org/doc/stable/reference/generated/numpy.take_along_axis.html
):
```
out_{ijk} = arr_{idx_{ijk}, j, k} if axis=0
              or arr_{i, idx_{ijk}, k} if axis=1
              or arr_{i, j, idx_{ijk}}  if axis=2
```
But `gather` is used to index across axis, including the preceding axis. This *conflicts* with the linear storage of `List<NDIndexElement>` inside `NDIndex`. So `gather` is treated parallel to `NDIndex` when fed into `get`.

Similar functions in numpy / pytorch / mxnet can also be implemented: e.g. `torch.take`, and advanced indexing like `arr[:, [0 ,2]]` and `arr[idx_r, idx_c]`

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

```diff
@@             Coverage Diff              @@
##             master    #1622      +/-   ##
============================================
- Coverage     72.08%   70.90%   -1.19%     
- Complexity     5126     5439     +313     
============================================
  Files           473      508      +35     
  Lines         21970    23773    +1803     
  Branches       2351     2588     +237     
============================================
+ Hits          15838    16856    +1018     
- Misses         4925     5628     +703     
- Partials       1207     1289      +82     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1622?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/1622/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: |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/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/1622/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/1622/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: |
| [.../modality/cv/translator/YoloTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/YoloV5Translator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvci5qYXZh) | `5.69% <0.00%> (ø)` | |
| [...odality/cv/translator/YoloV5TranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...pi/src/main/java/ai/djl/ndarray/BytesSupplier.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvQnl0ZXNTdXBwbGllci5qYXZh) | `54.54% <0.00%> (-12.13%)` | :arrow_down: |
| [...main/java/ai/djl/ndarray/index/NDArrayIndexer.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvaW5kZXgvTkRBcnJheUluZGV4ZXIuamF2YQ==) | `66.66% <ø> (ø)` | |
| [...ain/java/ai/djl/ndarray/index/dim/NDIndexPick.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvaW5kZXgvZGltL05ESW5kZXhQaWNrLmphdmE=) | `100.00% <ø> (ø)` | |
| ... and [245 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1622/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/1622?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/1622?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [bbdcbf7...d1fc3a9](https://codecov.io/gh/deepjavalibrary/djl/pull/1622?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 88eec5517..c35c7f095 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -545,6 +545,17 @@ public interface NDArray extends NDResource, BytesSupplier {
         return get(new NDIndex().addBooleanIndex(index));
     }
 
+    /**
+     * Returns a partial {@code NDArray} pointed by the indexed array. Given NDArray arr, NDArray
+     * idx, and long axis, the output is out_{ijk} = arr_{idx_{ijk}, j, k} if axis=0 or arr_{i,
+     * idx_{ijk}, k} if axis=1 or arr_{i, j, idx_{ijk}} if axis=2
+     *
+     * @param index picks the elements of an NDArray to the same position as index
+     * @param axis the entries of index are indices of axis
+     * @return the partial {@code NDArray} of the same shape as index
+     */
+    NDArray gather(NDArray index, int axis);
+
     /**
      * Returns a scalar {@code NDArray} corresponding to a single element.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 1a30edb19..46f2634c3 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -176,6 +176,12 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException(UNSUPPORTED_MSG);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gather(NDArray index, int axis) {
+        throw new UnsupportedOperationException(UNSUPPORTED_MSG);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void set(Buffer data) {
diff --git a/docs/README.md b/docs/README.md
index 3f46ce502..22b7815f6 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -4,6 +4,8 @@ This folder contains examples and documentation for the Deep Java Library (DJL)
 
 ## [JavaDoc API Reference](https://javadoc.djl.ai/)
 
+Note: when searching in JavaDoc, if your access is denied, please try removing the string `undefined` in the url.
+
 ## [Demos](https://djl.ai/website/demo.html)
 
 ## Cheat sheet
diff --git a/docs/get.md b/docs/get.md
index 39330e2f0..056c19445 100644
--- a/docs/get.md
+++ b/docs/get.md
@@ -90,9 +90,19 @@ If you build following the above instructions, you will use the version of the c
 
 You can look here to find the [list of DJL releases](https://github.com/deepjavalibrary/djl/releases).
 
-### Using build from source in another project
+### Using built-from-source version in another project
 
-If you have another project and want to use a custom version of DJL in it, you can do this by building from source. Once you have built DJL from source, run `./gradlew publishToMavenLocal`. This will install DJL to your local maven repository cache located on your filesystem at `~/.m2/repository`. When you publish here, it will use the same snapshot versions as used by the snapshot repository above for adding the DJL dependencies.
+If you have another project and want to use a custom version of DJL in it, you can do the following. First, build DJL from source by running `./gradlew build` inside djl folder. Then run `./gradlew publishToMavenLocal`, which will install DJL to your local maven repository cache, located on your filesystem at `~/.m2/repository`. After publishing it here, you can add the DJL snapshot version dependencies as shown below 
+
+```groovy
+dependencies {
+    implementation platform("ai.djl:bom:0.17.0-SNAPSHOT")
+}
+```
+
+This snapshot version is the same as the custom DJL repository. 
+
+You also need to change directory to `djl/bom`. Then build and publish it to maven local same as was done in `djl`.
 
 From there, you may have to update the Maven or Gradle build of the project importing DJL to also look at the local maven repository cache for your locally published versions of DJL. For Maven, no changes are necessary. If you are using Gradle, you will have to add the maven local repository such as this [example](https://github.com/deepjavalibrary/djl-demo/blob/135c969d66d98d1672852e53a37e52ca1da3e325/pneumonia-detection/build.gradle#L11):
 
@@ -101,3 +111,5 @@ repositories {
     mavenLocal()
 }
 ```
+
+Note that `mavenCentral()` may still be needed for applications like log4j and json.
\ No newline at end of file
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index 68ab1b61b..f786b0229 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -306,6 +306,12 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         JnaUtils.syncCopyFromCPU(getHandle(), buf, size);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gather(NDArray index, int axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void copyTo(NDArray ndArray) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index ebad11605..a1cacb83e 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -239,6 +239,15 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
         return JniUtils.getItem(this, indices);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gather(NDArray index, int axis) {
+        if (!(index instanceof PtNDArray)) {
+            throw new IllegalArgumentException("Only PtNDArray is supported.");
+        }
+        return JniUtils.gather(this, (PtNDArray) index, axis);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void copyTo(NDArray array) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 23fe2e276..a11719612 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
@@ -346,6 +346,12 @@ public final class JniUtils {
         PyTorchLibrary.LIB.torchSet(self.getHandle(), data);
     }
 
+    public static PtNDArray gather(PtNDArray ndArray, PtNDArray index, long dim) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchGather(ndArray.getHandle(), index.getHandle(), dim, false));
+    }
+
     public static PtNDArray pick(PtNDArray ndArray, PtNDArray index, long dim) {
         Shape indexShape = index.getShape();
         Shape ndShape = ndArray.getShape();
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 3d5c8f09c..026680b97 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
@@ -206,6 +206,12 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         JavacppUtils.setByteBuffer(getHandle(), buf);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray gather(NDArray index, int axis) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void attach(NDManager manager) {
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index f1ffe6612..fdd2345a0 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -3,4 +3,4 @@ distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.4.1-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.4.2-bin.zip


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 5bf2e6010..a954f3e0f 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
@@ -16,6 +16,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.index.NDIndex;
 import ai.djl.ndarray.types.Shape;
+import ai.djl.testing.TestRequirements;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -52,6 +53,22 @@ public class NDIndexTest {
         }
     }
 
+    @Test
+    public void testGather() {
+        // Currently in windows gradle cannot find all the engines.
+        // In the dependencies, changing runtimeOnly to api however will remedy the problem.
+        // TODO: remove this when gradle problem is fixed.
+        TestRequirements.notWindows();
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray arr = manager.arange(20f).reshape(-1, 4);
+            long[] idx = {0, 0, 2, 1, 1, 2};
+            NDArray index = manager.create(idx, new Shape(3, 2));
+            NDArray actual = arr.gather(index, 1);
+            NDArray expected = manager.create(new float[] {0, 0, 6, 5, 9, 10}, new Shape(3, 2));
+            Assert.assertEquals(actual, expected);
+        }
+    }
+
     @Test
     public void testGet() {
         try (NDManager manager = NDManager.newBaseManager()) {