split broadcastable op logic into simple elementwise method and broadcast method
@@ -1458,26 +1458,42 @@ public class JVMDoubleTensor extends DoubleTensor {
return opWithAutoBroadcast(this, that, op, true);
}
- private DoubleTensor opWithAutoBroadcast(DoubleTensor left, DoubleTensor right, BiFunction<Double, Double, Double> op, boolean inPlace) {
+ private JVMDoubleTensor opWithAutoBroadcast(DoubleTensor left, DoubleTensor right, BiFunction<Double, Double, Double> op, boolean inPlace) {
- double[] leftBuffer = left.asFlatDoubleArray();
- long[] leftShape = left.getShape();
- long[] leftStride = left.getStride();
+ final double[] leftBuffer = left.asFlatDoubleArray();
+ final long[] leftShape = left.getShape();
+ final long[] leftStride = left.getStride();
- double[] rightBuffer = right.asFlatDoubleArray();
- long[] rightShape = right.getShape();
- long[] rightStride = right.getStride();
+ final double[] rightBuffer = right.asFlatDoubleArray();
+ final long[] rightShape = right.getShape();
+ final long[] rightStride = right.getStride();
- boolean needsBroadcast = !Arrays.equals(leftShape, rightShape);
-
- long[] resultShape;
- double[] outputBuffer;
+ final boolean needsBroadcast = !Arrays.equals(leftShape, rightShape);
if (needsBroadcast) {
- resultShape = Shape.broadcastOutputShape(leftShape, rightShape);
+ return broadcastBinaryOp(
+ leftBuffer, leftShape, leftStride,
+ rightBuffer, rightShape, rightStride,
+ op, inPlace
+ );
+
+ } else {
+
+ return elementwiseBinaryOp(leftBuffer, rightBuffer, leftShape, op, inPlace);
+ }
+
+ }
+
+ private JVMDoubleTensor broadcastBinaryOp(double[] leftBuffer, long[] leftShape, long[] leftStride,
+ double[] rightBuffer, long[] rightShape, long[] rightStride,
+ BiFunction<Double, Double, Double> op,
+ boolean inPlace) {
+
+ long[] resultShape = Shape.broadcastOutputShape(leftShape, rightShape);
boolean resultShapeIsThisShape = Arrays.equals(resultShape, leftShape);
+ final double[] outputBuffer;
if (!resultShapeIsThisShape) {
boolean resultShapeIsThatShape = Arrays.equals(resultShape, rightShape);
@@ -1509,16 +1525,17 @@ public class JVMDoubleTensor extends DoubleTensor {
//e.g. [2] / [2, 2]
broadcastFromLeft(leftBuffer, leftShape, leftStride, rightBuffer, rightShape, rightStride, outputBuffer, op);
}
+ }
- } else {
- resultShape = leftShape;
+ private JVMDoubleTensor elementwiseBinaryOp(double[] leftBuffer, double[] rightBuffer, long[] shape,
+ BiFunction<Double, Double, Double> op,
+ boolean inPlace) {
- outputBuffer = inPlace ? leftBuffer : new double[leftBuffer.length];
+ final double[] outputBuffer = inPlace ? leftBuffer : new double[leftBuffer.length];
elementwise(leftBuffer, rightBuffer, outputBuffer, op);
- }
- return new JVMDoubleTensor(outputBuffer, resultShape);
+ return new JVMDoubleTensor(outputBuffer, shape);
}
@Override
