Repository: deepjavalibrary/djl
Pull request #2651 (2023-06-13 20:38:07)

[pytorch] Impoves IValue tuple of tuple support
## 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/2651?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Patch coverage: **`71.11`**% and project coverage change: **`+1.46`** :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 [(`9f14b96`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2651?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 73.55%.

: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    #2651      +/-   ##
============================================
+ Coverage     72.08%   73.55%   +1.46%     
- Complexity     5126     6973    +1847     
============================================
  Files           473      685     +212     
  Lines         21970    30451    +8481     
  Branches       2351     3157     +806     
============================================
+ Hits          15838    22399    +6561     
- Misses         4925     6522    +1597     
- Partials       1207     1530     +323     
```


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

... and [354 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2651/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/2651?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/jni/IValueUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java
index 7704acc5a..f1d977531 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/IValueUtils.java
@@ -26,6 +26,7 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 /** IValueUtils is utility class to deal with IValue in PyTorch. */
@@ -33,6 +34,7 @@ public final class IValueUtils {
 
     private static final Pattern PATTERN_LIST = Pattern.compile("\\w+\\[]");
     private static final Pattern PATTERN_TUPLE = Pattern.compile("\\w+\\(\\)");
+    private static final Pattern PATTERN_TUPLE_OF_TUPLE = Pattern.compile("\\w+(\\([\\d,]+\\))");
 
     private IValueUtils() {}
 
@@ -101,6 +103,7 @@ public final class IValueUtils {
         String methodName = "forward";
         for (NDArray array : ndList) {
             String name = array.getName();
+            Matcher m;
             if (name != null && name.contains(".")) {
                 String[] strings = name.split("\\.", 2);
                 int index = addToMap(indexMap, strings[0], outputs);
@@ -116,6 +119,11 @@ public final class IValueUtils {
                 int index = addToMap(indexMap, name, outputs);
                 PairList<String, PtNDArray> pl = outputs.get(index);
                 pl.add("()", (PtNDArray) array);
+            } else if (name != null && (m = PATTERN_TUPLE_OF_TUPLE.matcher(name)).matches()) {
+                int index = addToMap(indexMap, name, outputs);
+                String key = m.group(1);
+                PairList<String, PtNDArray> pl = outputs.get(index);
+                pl.add(key, (PtNDArray) array);
             } else {
                 PairList<String, PtNDArray> pl = new PairList<>();
                 pl.add(null, (PtNDArray) array);
@@ -137,6 +145,19 @@ public final class IValueUtils {
                 // Tuple
                 IValue[] arrays = pl.values().stream().map(IValue::from).toArray(IValue[]::new);
                 ret[i] = IValue.tupleFrom(arrays);
+            } else if (key.startsWith("(")) {
+                // Tuple of tuple
+                String[] keys = key.substring(1, key.length() - 1).split(",");
+                int[] dim = Arrays.stream(keys).mapToInt(Integer::parseInt).toArray();
+                List<PtNDArray> arrays = pl.values();
+                int product = 1;
+                for (int d : dim) {
+                    product *= d;
+                }
+                if (product != arrays.size()) {
+                    throw new IllegalArgumentException("Invalid NDList tuple size: " + key);
+                }
+                ret[i] = IValueUtils.toTupleIValueRecur(arrays, dim, 0, 0).getKey();
             } else {
                 Map<String, PtNDArray> map = new ConcurrentHashMap<>();
                 for (Pair<String, PtNDArray> pair : pl) {
@@ -148,47 +169,23 @@ public final class IValueUtils {
         return new Pair<>(ret, methodName);
     }
 
-    /**
-     * Converts ndList to IValue.
-     *
-     * @param ndList the NDList to convert
-     * @param dims the shape of the output
-     * @return the result {@link IValue}
-     */
-    public static IValue toTupleIValue(NDList ndList, long[] dims) {
-        return toTupleIValueRecur(ndList, dims, 0, 0).getKey();
-    }
-
-    /**
-     * Helper function.
-     *
-     * @param ndList the NDList to convert
-     * @param dims the shape of the output
-     * @param startCount the start index of the current recursion level
-     * @param level the recursion level
-     * @return the result
-     */
     private static Pair<IValue, Integer> toTupleIValueRecur(
-            NDList ndList, long[] dims, int startCount, int level) {
-        if (startCount > ndList.size()) {
-            throw new IllegalArgumentException("startCount illegal");
-        }
+            List<PtNDArray> list, int[] dims, int start, int level) {
         if (dims.length - 1 == level) {
-            long dim = dims[level];
-            List<PtNDArray> vector = new ArrayList<>();
-            for (int i = startCount; i < startCount + dim; i++) {
-                vector.add((PtNDArray) ndList.get(i));
+            int dim = dims[level];
+            IValue[] iValues = new IValue[dim];
+            for (int i = 0; i < dim; i++) {
+                iValues[i] = IValue.from(list.get(i + start));
             }
-            IValue[] output = vector.stream().map(IValue::from).toArray(IValue[]::new);
-            return new Pair<>(IValue.tupleFrom(output), Math.toIntExact((startCount + dim)));
+            return new Pair<>(IValue.tupleFrom(iValues), Math.toIntExact((start + dim)));
         }
 
-        IValue[] output = new IValue[Math.toIntExact(dims[0])];
+        IValue[] output = new IValue[dims[0]];
         for (int j = 0; j < dims[level]; j++) {
-            Pair<IValue, Integer> p = toTupleIValueRecur(ndList, dims, startCount, level + 1);
-            startCount = p.getValue();
+            Pair<IValue, Integer> p = toTupleIValueRecur(list, dims, start, level + 1);
+            start = p.getValue();
             output[j] = p.getKey();
         }
-        return new Pair<>(IValue.tupleFrom(output), startCount);
+        return new Pair<>(IValue.tupleFrom(output), start);
     }
 }


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
index 0d4769c84..273cae610 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/IValueUtilsTest.java
@@ -16,6 +16,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
+import ai.djl.pytorch.engine.PtNDManager;
 import ai.djl.util.Pair;
 
 import org.testng.Assert;
@@ -47,6 +48,47 @@ public class IValueUtilsTest {
         }
     }
 
+    @Test
+    public void testTupleOfTuple() {
+        try (PtNDManager manager = (PtNDManager) NDManager.newBaseManager()) {
+            NDArray array1 = manager.zeros(new Shape(1));
+            array1.setName("input1(2,3)");
+            NDArray array2 = manager.ones(new Shape(2));
+            array2.setName("input1(2,3)");
+            NDArray array3 = manager.ones(new Shape(3));
+            array3.setName("input1(2,3)");
+            NDArray array4 = manager.ones(new Shape(4));
+            array4.setName("input1(2,3)");
+            NDArray array5 = manager.ones(new Shape(5));
+            array5.setName("input1(2,3)");
+            NDArray array6 = manager.ones(new Shape(6));
+            array6.setName("input1(2,3)");
+            NDList input = new NDList(array1, array2, array3, array4, array5, array6);
+            // the NDList is mapped to (input1: Tuple((t1, t2, t3), (t34, t5, t6))
+            input.attach(manager);
+
+            Pair<IValue[], String> result = IValueUtils.getInputs(input);
+            IValue[] iValues = result.getKey();
+            Assert.assertEquals(result.getValue(), "forward");
+            Assert.assertEquals(iValues.length, 1);
+            Assert.assertTrue(iValues[0].isTuple());
+
+            IValue[] tuple = iValues[0].toIValueTuple();
+            Assert.assertEquals(tuple.length, 2);
+            Assert.assertTrue(tuple[0].isTuple());
+            IValue[] subTuple = tuple[1].toIValueTuple();
+            Assert.assertEquals(subTuple.length, 3);
+            NDList output = iValues[0].toNDList(manager);
+            Assert.assertEquals(output.size(), 6);
+            Assert.assertEquals(output.get(5).getShape().get(0), 6);
+
+            Arrays.stream(iValues).forEach(IValue::close);
+
+            NDList input2 = new NDList(array1, array2);
+            Assert.assertThrows(() -> IValueUtils.getInputs(input2));
+        }
+    }
+
     @Test
     public void testMapOfTensor() {
         try (NDManager manager = NDManager.newBaseManager()) {