refactor pluck dual
@@ -322,16 +322,22 @@ public class DualNumber {
public DualNumber pluck(int[] inputShape, int... index) {
Map<Long, DoubleTensor> pluckedDuals = new HashMap<>();
-
long inputLength = TensorShape.getLength(inputShape);
int flatIndex = TensorShape.getFlatIndex(inputShape, TensorShape.getRowFirstStride(inputShape), index);
for (Map.Entry<Long, DoubleTensor> entry : this.partialDerivatives.asMap().entrySet()) {
- DoubleTensor partial = entry.getValue();
- int[] partialShape = entry.getValue().getShape();
+ DoubleTensor plucked = pluckedFromPartial(entry.getValue(), inputLength, flatIndex);
+ pluckedDuals.put(entry.getKey(), plucked);
+ }
+
+ return new DualNumber(DoubleTensor.scalar(this.value.getValue(index)), pluckedDuals);
+ }
+
+ private DoubleTensor pluckedFromPartial(DoubleTensor partial, long inputLength, int inputFlatIndex) {
+ int[] partialShape = partial.getShape();
long partialLength = TensorShape.getLength(partialShape);
long howManyValuesToPluckFromPartial = partialLength / inputLength;
- int flatIndexOfPartial = (int) howManyValuesToPluckFromPartial * flatIndex;
+ int flatIndexOfPartial = (int) howManyValuesToPluckFromPartial * inputFlatIndex;
double[] flatPartial = partial.asFlatDoubleArray();
double[] pluckedValues = new double[(int) howManyValuesToPluckFromPartial];
@@ -341,13 +347,7 @@ public class DualNumber {
int[] newShape = Arrays.copyOfRange(partialShape, partialShape.length - 2, partialShape.length);
newShape = TensorShape.shapeToDesiredRankByPrependingOnes(newShape, partial.getShape().length);
- DoubleTensor plucked = DoubleTensor.create(pluckedValues, newShape);
- pluckedDuals.put(entry.getKey(), plucked);
-
- }
-
- return new DualNumber(DoubleTensor.scalar(this.value.getValue(index)), pluckedDuals);
+ return DoubleTensor.create(pluckedValues, newShape);
}
-
}
