work around for potential nd4j bug
package io.improbable.keanu.tensor.dbl;
+import com.google.common.primitives.Ints;
import io.improbable.keanu.tensor.Tensor;
import io.improbable.keanu.tensor.bool.BooleanTensor;
import io.improbable.keanu.tensor.bool.SimpleBooleanTensor;
@@ -13,12 +14,15 @@ import org.nd4j.linalg.api.buffer.DataBuffer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.ops.impl.broadcast.BroadcastMulOp;
import org.nd4j.linalg.api.ops.impl.transforms.comparison.*;
+import org.nd4j.linalg.api.shape.Shape;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.conditions.Conditions;
import org.nd4j.linalg.inverse.InvertMatrix;
import org.nd4j.linalg.ops.transforms.Transforms;
+import java.util.ArrayList;
import java.util.Arrays;
+import java.util.List;
import java.util.function.Function;
import static java.util.Arrays.copyOf;
@@ -341,19 +345,37 @@ public class Nd4jDoubleTensor implements DoubleTensor {
return that.times(this.scalar());
} else {
INDArray thatArray = unsafeGetNd4J(that);
-// if (Arrays.equals(tensor.shape(), thatArray.shape())) {
+ if (Arrays.equals(tensor.shape(), thatArray.shape())) {
return new Nd4jDoubleTensor(tensor.mul(thatArray));
-// } else {
-// INDArray result = Nd4j.createUninitialized(tensor.shape(), tensor.ordering());
-// int[] broadcastDimensions = new int[]{0};
-// Nd4j.getExecutioner().exec(
-// new BroadcastMulOp(tensor, thatArray, result, broadcastDimensions),
-// broadcastDimensions
-// );
-//
-// return new Nd4jDoubleTensor(result);
-// }
+ } else {
+ int[] broadcastDimensions = Shape.getBroadcastDimensions(tensor.shape(), thatArray.shape());
+ int[] executeAlong = getBroadcastAlongDimensions(tensor.shape(), thatArray.shape());
+
+ INDArray result = broadcastMultiply(tensor, thatArray, broadcastDimensions, executeAlong);
+
+ return new Nd4jDoubleTensor(result);
+ }
+ }
+ }
+
+ private static INDArray broadcastMultiply(INDArray a, INDArray b, int[] broadcastDimensions, int[] alongDimensions) {
+ INDArray result = Nd4j.createUninitialized(a.shape(), a.ordering());
+ Nd4j.getExecutioner().exec(
+ new BroadcastMulOp(a, b, result, broadcastDimensions),
+ alongDimensions
+ );
+ return result;
+ }
+
+ private static int[] getBroadcastAlongDimensions(int[] shapeA, int[] shapeB) {
+ int minRank = Math.min(shapeA.length, shapeB.length);
+ List<Integer> along = new ArrayList<>();
+ for (int i = 0; i < minRank; i++) {
+ if (shapeA[i] == shapeB[i]) {
+ along.add(i);
+ }
}
+ return Ints.toArray(along);
}
@Override
