remove duplicate correct for scalar broadcast check and add javadoc
@@ -13,7 +13,7 @@ public class AutoDiffBroadcast {
public static PartialDerivative correctForScalarPartialForward(PartialDerivative partial, long[] partialOfShape, long[] targetOfShape) {
- if (shouldCorrectPartialForScalarForward(partial, partialOfShape, targetOfShape)) {
+ if (shouldCorrectPartialForScalar(partial, partialOfShape, targetOfShape)) {
long[] wrtShape = partial.getWrtShape(partialOfShape);
DoubleTensor correctedPartial = DoubleTensor
@@ -26,13 +26,9 @@ public class AutoDiffBroadcast {
}
}
- private static boolean shouldCorrectPartialForScalarForward(PartialDerivative partial, long[] partialOfShape, long[] targetOfShape) {
- return partial.isPresent() && !Arrays.equals(partialOfShape, targetOfShape);
- }
-
public static PartialDerivative correctForScalarPartialReverse(PartialDerivative partial, long[] partialWrtShape, long[] targetWrtShape) {
- if (shouldCorrectForPartialScalarReverse(partial, partialWrtShape, targetWrtShape)) {
+ if (shouldCorrectPartialForScalar(partial, partialWrtShape, targetWrtShape)) {
int[] wrtDims = TensorShape.dimensionRange(-partialWrtShape.length, 0);
DoubleTensor partialSummed = partial.getPartial().sum(wrtDims);
@@ -48,7 +44,15 @@ public class AutoDiffBroadcast {
}
}
- public static boolean shouldCorrectForPartialScalarReverse(PartialDerivative partial, long[] partialWrtShape, long[] targetWrtShape) {
- return partial.isPresent() && !Arrays.equals(partialWrtShape, targetWrtShape);
+ /**
+ * @param partial The partial derivative that may or may not come from a broadcasted operation.
+ * @param actualShape The part of the partial shape that should match the expected shape in the case no broadcast was
+ * performed. This would be the of shape for forward mode and the with respect to shape for reverse.
+ * @param expectedShape The shape of the operation result in forward mode or the shape of the operand in reverse mode.
+ * This should match the actual shape if no broadcast was performed.
+ * @return true if a broadcast should be taken into account and corrected for in the auto diff calculation, false otherwise.
+ */
+ private static boolean shouldCorrectPartialForScalar(PartialDerivative partial, long[] actualShape, long[] expectedShape) {
+ return partial.isPresent() && !Arrays.equals(actualShape, expectedShape);
}
}
