@hoff97/tensor-js/dist/lib/autograd/ops/reduce/sumBack.js

PyTorch like deep learning inferrence library

export class SumBack {
    constructor(input, sumDims, keepDims) {
        this.input = input;
        this.sumDims = sumDims;
        this.keepDims = keepDims;
    }
    backward(grad) {
        const inShape = this.input.value.getShape();
        if (!this.keepDims) {
            const newShape = [];
            let sumI = 0;
            for (let i = 0; i < inShape.length; i++) {
                if (sumI < this.sumDims.length && this.sumDims[sumI] === i) {
                    newShape.push(1);
                    sumI++;
                }
                else {
                    newShape.push(inShape[i]);
                }
            }
            grad = grad.reshape(newShape, false);
        }
        grad = grad.expand(inShape);
        const needed = this.input.backward(grad);
        if (!needed) {
            grad.delete();
        }
    }
    delete() {
        if (!this.input.isLeaf()) {
            this.input.delete();
        }
    }
}
//# sourceMappingURL=sumBack.js.map