[fix]: bugfix choice(INDArray source, INDArray probs, int numSamples)

deeplearning4j · Jun 9, 2021 · ba60684 · ba60684
1 parent 377e19c
commit ba60684
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/factory/Nd4j.java b/nd4j/nd4j-backends/nd4j-api-parent/nd4j-api/src/main/java/org/nd4j/linalg/factory/Nd4j.java
@@ -2758,7 +2758,7 @@ public static INDArray choice(INDArray source, INDArray probs, int numSamples,
         if (numSamples < 1)
             throw new ND4JIllegalStateException("Nd4j.choice() numSamples must be positive value");
 
-        return choice(source, probs, createUninitialized(numSamples), rng);
+        return choice(source, probs, createUninitialized(source.dataType(), numSamples), rng);
     }
 
     // @see tag works well here.

diff --git a/nd4j/nd4j-backends/nd4j-tests/src/test/java/org/nd4j/linalg/factory/Nd4jTest.java b/nd4j/nd4j-backends/nd4j-tests/src/test/java/org/nd4j/linalg/factory/Nd4jTest.java
@@ -289,6 +289,17 @@ public void testNpyByteArray() throws Exception {
 
     }
 
+    @Test
+    public void testChoiceDataType() {
+        INDArray dataTypeIsDouble = Nd4j.empty(DataType.DOUBLE);
+
+        INDArray source = Nd4j.createFromArray(new double[] { 1.0, 0.0 });
+        INDArray probs = Nd4j.valueArrayOf(new long[] { 2 }, 0.5, DataType.DOUBLE);
+        INDArray actual = Nd4j.choice(s, p, 10);
+
+
+        assertEquals(dataTypeIsDouble.dataType(), actual.dataType());
+    }
 
 
 }