Repository: deepjavalibrary/djl
Pull request #3287 (2024-06-28 17:41:21)

[mxnet] Fixes GloveWordEmbeddingTranslator bug
Fixes: #3286

## 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



Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index 4cc020ea0..8f26e2713 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -399,7 +399,7 @@ public interface NDManager extends AutoCloseable {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(float[][] data) {
-        FloatBuffer buffer = FloatBuffer.allocate(data.length * data[0].length);
+        FloatBuffer buffer = allocateDirect(data.length * data[0].length * 4).asFloatBuffer();
         for (float[] d : data) {
             buffer.put(d);
         }
@@ -414,7 +414,7 @@ public interface NDManager extends AutoCloseable {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(int[][] data) {
-        IntBuffer buffer = IntBuffer.allocate(data.length * data[0].length);
+        IntBuffer buffer = allocateDirect(data.length * data[0].length * 4).asIntBuffer();
         for (int[] d : data) {
             buffer.put(d);
         }
@@ -429,7 +429,7 @@ public interface NDManager extends AutoCloseable {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(double[][] data) {
-        DoubleBuffer buffer = DoubleBuffer.allocate(data.length * data[0].length);
+        DoubleBuffer buffer = allocateDirect(data.length * data[0].length * 8).asDoubleBuffer();
         for (double[] d : data) {
             buffer.put(d);
         }
@@ -444,7 +444,7 @@ public interface NDManager extends AutoCloseable {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(long[][] data) {
-        LongBuffer buffer = LongBuffer.allocate(data.length * data[0].length);
+        LongBuffer buffer = allocateDirect(data.length * data[0].length * 8).asLongBuffer();
         for (long[] d : data) {
             buffer.put(d);
         }
@@ -459,7 +459,7 @@ public interface NDManager extends AutoCloseable {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(byte[][] data) {
-        ByteBuffer buffer = ByteBuffer.allocate(data.length * data[0].length);
+        ByteBuffer buffer = allocateDirect(data.length * data[0].length);
         for (byte[] d : data) {
             buffer.put(d);
         }
@@ -474,7 +474,7 @@ public interface NDManager extends AutoCloseable {
      * @return a new instance of {@link NDArray}
      */
     default NDArray create(boolean[][] data) {
-        ByteBuffer buffer = ByteBuffer.allocate(data.length * data[0].length);
+        ByteBuffer buffer = allocateDirect(data.length * data[0].length);
         for (boolean[] d : data) {
             for (boolean b : d) {
                 buffer.put((byte) (b ? 1 : 0));
diff --git a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java
index de9cb4ffc..a5d6213c1 100644
--- a/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java
+++ b/engines/mxnet/mxnet-model-zoo/src/main/java/ai/djl/mxnet/zoo/nlp/embedding/GloveWordEmbeddingTranslatorFactory.java
@@ -70,6 +70,7 @@ public class GloveWordEmbeddingTranslatorFactory implements TranslatorFactory {
         /** {@inheritDoc} */
         @Override
         public NDList processOutput(TranslatorContext ctx, NDList list) {
+            list.detach();
             return list;
         }
 


Test patch:
diff --git a/extensions/timeseries/src/test/java/ai/djl/timeseries/distribution/DistributionTest.java b/extensions/timeseries/src/test/java/ai/djl/timeseries/distribution/DistributionTest.java
index ecb470a2a..47ad58e7b 100644
--- a/extensions/timeseries/src/test/java/ai/djl/timeseries/distribution/DistributionTest.java
+++ b/extensions/timeseries/src/test/java/ai/djl/timeseries/distribution/DistributionTest.java
@@ -37,15 +37,17 @@ public class DistributionTest {
             NDArray logits = manager.create(new float[] {0.1f, 2f});
             totalCount.setName("total_count");
             logits.setName("logits");
-            Distribution negativeBinomial =
+            Distribution negativeBinomialDistribution =
                     NegativeBinomial.builder().setDistrArgs(new NDList(totalCount, logits)).build();
 
-            NDArray expected = manager.create(new float[] {-2.3027f, -2.2539f});
-            NDArray real = negativeBinomial.logProb(manager.create(new float[] {2f, 1f}));
-            Assertions.assertAlmostEquals(real, expected);
+            NDArray expectedValues = manager.create(new float[] {-2.3027f, -2.2539f});
+            NDArray actualValues =
+                    negativeBinomialDistribution.logProb(manager.create(new float[] {2f, 1f}));
+            Assertions.assertAlmostEquals(actualValues, expectedValues);
 
-            NDArray samplesMean = negativeBinomial.sample(100000).mean(new int[] {0});
-            Assertions.assertAlmostEquals(samplesMean, negativeBinomial.mean(), 2e-2f, 2e-2f);
+            NDArray samplesMean = negativeBinomialDistribution.sample(100000).mean(new int[] {0});
+            Assertions.assertAlmostEquals(
+                    samplesMean, negativeBinomialDistribution.mean(), 2e-2f, 2e-2f);
         }
     }
 
@@ -58,15 +60,16 @@ public class DistributionTest {
             mu.setName("mu");
             sigma.setName("sigma");
             nu.setName("nu");
-            Distribution studentT =
+            Distribution studentTDistribution =
                     StudentT.builder().setDistrArgs(new NDList(mu, sigma, nu)).build();
 
-            NDArray expected = manager.create(new float[] {-0.9779f, -1.6940f});
-            NDArray real = studentT.logProb(manager.create(new float[] {1000f, -1000f}));
-            Assertions.assertAlmostEquals(real, expected);
+            NDArray expectedValues = manager.create(new float[] {-0.9779f, -1.6940f});
+            NDArray actualValues =
+                    studentTDistribution.logProb(manager.create(new float[] {1000f, -1000f}));
+            Assertions.assertAlmostEquals(actualValues, expectedValues);
 
-            NDArray samplesMean = studentT.sample(100000).mean(new int[] {0});
-            Assertions.assertAlmostEquals(samplesMean, studentT.mean(), 2e-2f, 2e-2f);
+            NDArray samplesMean = studentTDistribution.sample(100000).mean(new int[] {0});
+            Assertions.assertAlmostEquals(samplesMean, studentTDistribution.mean(), 2e-2f, 2e-2f);
         }
     }
 }