Repository: deepjavalibrary/djl
Pull request #2701 (2023-07-10 21:54:39)

[Pytorch] add basic median support
## Description ##

Add PyTorch median operator implementation, not on the frontend

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Patch coverage: **`54.72`**% and project coverage change: **`-0.13`** :warning:
> Comparison is base [(`bb5073f`)](https://app.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) 72.08% compared to head [(`4a2591c`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 71.96%.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2701      +/-   ##
============================================
- Coverage     72.08%   71.96%   -0.13%     
- Complexity     5126     7009    +1883     
============================================
  Files           473      698     +225     
  Lines         21970    31258    +9288     
  Branches       2351     3225     +874     
============================================
+ Hits          15838    22496    +6658     
- Misses         4925     7225    +2300     
- Partials       1207     1537     +330     
```


| [Impacted Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?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://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?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://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...djl/modality/cv/transform/RandomFlipLeftRight.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21GbGlwTGVmdFJpZ2h0LmphdmE=) | `25.00% <0.00%> (-25.00%)` | :arrow_down: |
| [...djl/modality/cv/transform/RandomFlipTopBottom.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21GbGlwVG9wQm90dG9tLmphdmE=) | `25.00% <0.00%> (-25.00%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...ain/java/ai/djl/modality/cv/util/NDImageUtils.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3V0aWwvTkRJbWFnZVV0aWxzLmphdmE=) | `67.10% <0.00%> (+7.89%)` | :arrow_up: |
| [api/src/main/java/ai/djl/modality/nlp/Decoder.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9EZWNvZGVyLmphdmE=) | `63.63% <ø> (ø)` | |
| ... and [226 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [368 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2701?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
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 14195d17b..e1f75ff4b 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
@@ -1482,13 +1482,19 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     /** {@inheritDoc} */
     @Override
     public PtNDArray median() {
-        throw new UnsupportedOperationException("Not implemented");
+        return median(new int[] {-1});
     }
 
     /** {@inheritDoc} */
     @Override
     public PtNDArray median(int[] axes) {
-        throw new UnsupportedOperationException("Not implemented");
+        if (axes.length != 1) {
+            throw new UnsupportedOperationException(
+                    "Not supporting zero or multi-dimension median");
+        }
+        NDList result = JniUtils.median(this, axes[0], false);
+        result.get(1).close();
+        return (PtNDArray) result.get(0);
     }
 
     /** {@inheritDoc} */
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 694387f33..aad38ae8f 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
@@ -875,6 +875,13 @@ public final class JniUtils {
                 PyTorchLibrary.LIB.torchMin(ndArray.getHandle(), dim, keepDim));
     }
 
+    public static NDList median(PtNDArray ndArray, long dim, boolean keepDim) {
+        long[] handles = PyTorchLibrary.LIB.torchMedian(ndArray.getHandle(), dim, keepDim);
+        return new NDList(
+                new PtNDArray(ndArray.getManager(), handles[0]),
+                new PtNDArray(ndArray.getManager(), handles[1]));
+    }
+
     public static PtNDArray mean(PtNDArray ndArray) {
         return new PtNDArray(
                 ndArray.getManager(), PyTorchLibrary.LIB.torchMean(ndArray.getHandle()));
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 65f5e8a47..63524ff60 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
@@ -238,6 +238,8 @@ final class PyTorchLibrary {
 
     native long torchMinimum(long self, long other);
 
+    native long[] torchMedian(long self, long dim, boolean keepDim);
+
     native long torchMin(long handle);
 
     native long torchMin(long handle, long dim, boolean keepDim);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
index c30d1e8eb..cfc4e9768 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
@@ -213,6 +213,20 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchMinimum(
   API_END_RETURN()
 }
 
+JNIEXPORT jlongArray JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchMedian(
+        JNIEnv* env, jobject jthis, jlong jself, jlong jdim, jboolean keep_dim) {
+    API_BEGIN()
+        const auto* self_ptr = reinterpret_cast<torch::Tensor*>(jself);
+        const auto result = self_ptr->median(jdim, keep_dim);
+        const auto* value_ptr = new torch::Tensor(std::get<0>(result));
+        const auto* indices_ptr = new torch::Tensor(std::get<1>(result));
+        std::vector<uintptr_t> vect;
+        vect.push_back(reinterpret_cast<uintptr_t>(value_ptr));
+        vect.push_back(reinterpret_cast<uintptr_t>(indices_ptr));
+        return djl::utils::jni::GetLongArrayFromVec(env, vect);
+    API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchAbs(JNIEnv* env, jobject jthis, jlong jhandle) {
   API_BEGIN()
   const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementComparisonOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementComparisonOpTest.java
index 7d4da81c7..1f5b215bc 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementComparisonOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementComparisonOpTest.java
@@ -601,6 +601,16 @@ public class NDArrayElementComparisonOpTest {
         }
     }
 
+    @Test
+    public void testMedian() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array1 = manager.create(new float[] {1, 3, 2, 5, 4});
+            Assert.assertEquals(array1.median(), manager.create(3.0f));
+            array1 = manager.create(new float[] {1, 3, 2, 5, 4, 8}, new Shape(2, 3));
+            Assert.assertEquals(array1.median(new int[] {1}), manager.create(new float[] {2, 5}));
+        }
+    }
+
     @Test
     public void testWhere() {
         try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {