Repository: deepjavalibrary/djl
Pull request #2780 (2023-09-13 21:49:09)

[pytorch] Avoid toByteBuffer() crash for large tensor
## 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://app.codecov.io/gh/deepjavalibrary/djl/pull/2780?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Patch coverage: **`55.15%`** and project coverage change: **`+0.02%`** :tada:
> 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 [(`fd61db3`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2780?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.11%.
> Report is 881 commits behind head on master.

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

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


```diff
@@             Coverage Diff              @@
##             master    #2780      +/-   ##
============================================
+ Coverage     72.08%   72.11%   +0.02%     
- Complexity     5126     7080    +1954     
============================================
  Files           473      700     +227     
  Lines         21970    31534    +9564     
  Branches       2351     3262     +911     
============================================
+ Hits          15838    22740    +6902     
- Misses         4925     7241    +2316     
- Partials       1207     1553     +346     
```


| [Files Changed](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2780?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/2780?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/2780?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/2780?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/2780?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/2780?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/2780?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/2780?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/2780?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/2780?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/2780?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/2780?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [370 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2780/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/2780?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](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-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 f33f1061b..e6aa2c145 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
@@ -308,8 +308,13 @@ JNIEXPORT jbyteArray JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchDataPtr
   // sparse and mkldnn are required to be converted to dense to access data ptr
   auto tensor = (tensor_ptr->is_sparse() || tensor_ptr->is_mkldnn()) ? tensor_ptr->to_dense() : *tensor_ptr;
   tensor = (tensor.is_contiguous()) ? tensor : tensor.contiguous();
-  jbyteArray result = env->NewByteArray(tensor.nbytes());
-  env->SetByteArrayRegion(result, 0, tensor.nbytes(), static_cast<const jbyte*>(tensor.data_ptr()));
+  size_t nbytes = tensor.nbytes();
+  if (nbytes > 0x7fffffff) {
+    env->ThrowNew(ENGINE_EXCEPTION_CLASS, "toByteBuffer() is not supported for large tensor");
+    return env->NewByteArray(0);
+  }
+  jbyteArray result = env->NewByteArray(nbytes);
+  env->SetByteArrayRegion(result, 0, nbytes, static_cast<const jbyte*>(tensor.data_ptr()));
   return result;
   API_END_RETURN()
 }


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/PtNDArrayTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/PtNDArrayTest.java
index 0423725a6..83b01f7a1 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/PtNDArrayTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/PtNDArrayTest.java
@@ -12,8 +12,11 @@
  */
 package ai.djl.pytorch.integration;
 
+import ai.djl.engine.EngineException;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
 
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -34,4 +37,12 @@ public class PtNDArrayTest {
             Assert.assertThrows(UnsupportedOperationException.class, () -> arr.get(0));
         }
     }
+
+    @Test
+    public void testLargeTensor() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(10 * 2850, 18944), DataType.FLOAT32);
+            Assert.assertThrows(EngineException.class, array::toByteArray);
+        }
+    }
 }