@tensorflow/tfjs-core/dist/tape_test.js

/**
 * @license
 * Copyright 2017 Google LLC. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */
import { add } from './engine';
import * as tf from './index';
import { ALL_ENVS, describeWithFlags } from './jasmine_util';
import { zerosLike } from './ops/ops';
import { backpropagateGradients, getFilteredNodesXToY } from './tape';
import { expectArraysClose } from './test_util';
describeWithFlags('getFilteredNodesXToY', ALL_ENVS, () => {
    it('no paths from x to y', () => {
        const x = tf.scalar(1);
        const intermediate1 = tf.scalar(0);
        const intermediate2 = tf.scalar(0);
        const y = tf.scalar(2);
        const tape = [
            {
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [intermediate1],
                gradient: null
            },
            {
                id: 1,
                kernelName: 'node1',
                inputs: { intermediate2 },
                outputs: [y],
                gradient: null
            }
        ];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
        expect(filteredTapeNodes.length).toBe(0);
        expect(filteredTapeNodes).toEqual([]);
    });
    it('one operation x => y', () => {
        const x = tf.scalar(1);
        const y = tf.scalar(2);
        const tape = [
            { id: 0, kernelName: 'node0', inputs: { x }, outputs: [y], gradient: null }
        ];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
        expect(filteredTapeNodes.length).toBe(1);
        expect(filteredTapeNodes).toEqual(tape);
    });
    it('1 operation [x0, x1] => y, all input paths', () => {
        const x0 = tf.scalar(0);
        const x1 = tf.scalar(1);
        const y = tf.scalar(2);
        const tape = [{
                id: 0,
                kernelName: 'node0',
                inputs: { x0, x1 },
                outputs: [y],
                gradient: null
            }];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x0, x1], y);
        expect(filteredTapeNodes.length).toBe(1);
        expect(filteredTapeNodes).toEqual(tape);
    });
    it('one operation [x0, x1] => y, one input paths', () => {
        const x0 = tf.scalar(0);
        const x1 = tf.scalar(1);
        const y = tf.scalar(2);
        const tape = [{
                id: 0,
                kernelName: 'node0',
                inputs: { x0, x1 },
                outputs: [y],
                gradient: null
            }];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x0], y);
        expect(filteredTapeNodes.length).toBe(1);
        // x1 input should be pruned, we don't ask for the gradient of x1.
        expect(filteredTapeNodes[0]).toEqual({
            id: 0,
            kernelName: 'node0',
            inputs: { x0 },
            outputs: [y],
            gradient: null
        });
    });
    it('two operations x => intermediate => y', () => {
        const x = tf.scalar(1);
        const intermediate = tf.scalar(0);
        const y = tf.scalar(2);
        const tape = [
            {
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [intermediate],
                gradient: null
            },
            {
                id: 1,
                kernelName: 'node1',
                inputs: { intermediate },
                outputs: [y],
                gradient: null
            }
        ];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
        expect(filteredTapeNodes.length).toBe(2);
        expect(filteredTapeNodes).toEqual(tape);
    });
    it('two operations [x0, x1], [x2] => ' +
        'intermediate => y', () => {
        const x0 = tf.scalar(1);
        const x1 = tf.scalar(2);
        const x2 = tf.scalar(3);
        const intermediate = tf.scalar(4);
        const y = tf.scalar(2);
        const tape = [
            {
                id: 0,
                kernelName: 'node0',
                inputs: { x0, x1 },
                outputs: [intermediate],
                gradient: null
            },
            {
                id: 1,
                kernelName: 'node1',
                inputs: { x2, intermediate },
                outputs: [y],
                gradient: null
            }
        ];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x0, x1, x2], y);
        expect(filteredTapeNodes.length).toBe(2);
        expect(filteredTapeNodes).toEqual(tape);
    });
    it('x => y and x => orphan', () => {
        const x = tf.scalar(1);
        const orphan = tf.scalar(0);
        const y = tf.scalar(2);
        const tape = [
            {
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [orphan],
                gradient: null
            },
            { id: 1, kernelName: 'node1', inputs: { x }, outputs: [y], gradient: null }
        ];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
        expect(filteredTapeNodes.length).toBe(1);
        // The orphan should be removed.
        expect(filteredTapeNodes[0]).toEqual(tape[1]);
    });
    it('x => y and orphan => y', () => {
        const x = tf.scalar(1);
        const orphan = tf.scalar(0);
        const y = tf.scalar(2);
        const tape = [{
                id: 0,
                kernelName: 'node0',
                inputs: { x, orphan },
                outputs: [y],
                gradient: null
            }];
        const filteredTapeNodes = getFilteredNodesXToY(tape, [x], y);
        expect(filteredTapeNodes.length).toBe(1);
        // The orphan should be pruned from the node's input.
        expect(filteredTapeNodes[0]).toEqual({
            id: 0,
            kernelName: 'node0',
            inputs: { x },
            outputs: [y],
            gradient: null
        });
    });
    it('1 op with 3 outputs x => y1, y2, y3', () => {
        const x = tf.scalar(1);
        const y1 = tf.scalar(2);
        const y2 = tf.scalar(2);
        const y3 = tf.scalar(2);
        const tape = [{
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [y1, y2, y3],
                gradient: null
            }];
        const filteredNodes1 = getFilteredNodesXToY(tape, [x], y1);
        expect(filteredNodes1.length).toBe(1);
        expect(filteredNodes1).toEqual(tape);
        const filteredNodes2 = getFilteredNodesXToY(tape, [x], y2);
        expect(filteredNodes2.length).toBe(1);
        expect(filteredNodes2).toEqual(tape);
        const filteredNodes3 = getFilteredNodesXToY(tape, [x], y3);
        expect(filteredNodes3.length).toBe(1);
        expect(filteredNodes3).toEqual(tape);
    });
});
describeWithFlags('backpropagateGradients', ALL_ENVS, () => {
    it('Throws if gradient is not defined', () => {
        const x = tf.scalar(0);
        const y = tf.scalar(1);
        const dy = tf.scalar(1);
        const accumulatedGradientsMap = {};
        accumulatedGradientsMap[y.id] = dy;
        const tape = [
            { id: 0, kernelName: 'node0', inputs: { x }, outputs: [y], gradient: null }
        ];
        expect(() => backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add))
            .toThrowError();
    });
    it('basic backprop with 1 node', async () => {
        const x = tf.scalar(0);
        const y = tf.scalar(1);
        const dy = tf.scalar(1);
        const accumulatedGradientsMap = {};
        accumulatedGradientsMap[y.id] = dy;
        const tape = [{
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [y],
                gradient: (dys) => {
                    return { x: () => dys[0].add(tf.scalar(1)) };
                }
            }];
        backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
        expectArraysClose(await accumulatedGradientsMap[x.id].data(), [2]);
    });
    it('basic backprop with 2 nodes', async () => {
        const x = tf.scalar(0);
        const intermediate = tf.scalar(1);
        const y = tf.scalar(2);
        const dy = tf.scalar(1);
        const accumulatedGradientsMap = {};
        accumulatedGradientsMap[y.id] = dy;
        const tape = [
            {
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [intermediate],
                gradient: (dys) => {
                    return { x: () => dys[0].add(tf.scalar(1)) };
                }
            },
            {
                id: 1,
                kernelName: 'node1',
                inputs: { intermediate },
                outputs: [y],
                gradient: (dys) => {
                    return { intermediate: () => dys[0].add(tf.scalar(1)) };
                }
            }
        ];
        backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
        // dx = dy + 1 + 1
        expectArraysClose(await accumulatedGradientsMap[x.id].data(), [3]);
    });
    it('basic backprop with a split node accumulates gradients', async () => {
        const x = tf.scalar(0);
        const intermediate1 = tf.scalar(1);
        const intermediate2 = tf.scalar(2);
        const y = tf.scalar(3);
        const dy = tf.scalar(1);
        const accumulatedGradientsMap = {};
        accumulatedGradientsMap[y.id] = dy;
        const tape = [
            {
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [intermediate1],
                gradient: (dys) => {
                    return { x: () => dys[0].add(tf.scalar(1)) };
                }
            },
            {
                id: 1,
                kernelName: 'node1',
                inputs: { x },
                outputs: [intermediate2],
                gradient: (dys) => {
                    return { x: () => dys[0].add(tf.scalar(1)) };
                }
            },
            {
                id: 2,
                kernelName: 'node2',
                inputs: { intermediate1, intermediate2 },
                outputs: [y],
                gradient: (dys) => {
                    return {
                        intermediate1: () => dys[0].add(tf.scalar(1)),
                        intermediate2: () => dys[0].add(tf.scalar(1))
                    };
                }
            }
        ];
        backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
        // dx = dy + 1 + 1 + 1 + 1 + 1
        expectArraysClose(await accumulatedGradientsMap[x.id].data(), [(await dy.data())[0] + 5]);
    });
    it('backprop over 1 node with 3 outputs, w.r.t to the 2nd output', async () => {
        const x = tf.tensor1d([1, 1, 1]);
        const y1 = tf.scalar(1);
        const y2 = tf.scalar(1);
        const y3 = tf.scalar(1);
        const accumulatedGradientsMap = {};
        // Backproping through the 2nd output.
        const dy2 = tf.scalar(5);
        accumulatedGradientsMap[y2.id] = dy2;
        let dys;
        const tape = [{
                id: 0,
                kernelName: 'node0',
                inputs: { x },
                outputs: [y1, y2, y3],
                gradient: (dys_) => {
                    dys = dys_.map(dy => dy || zerosLike(y1));
                    return { x: () => tf.stack(dys) };
                }
            }];
        backpropagateGradients(accumulatedGradientsMap, tape, f => tf.tidy(f), add);
        expectArraysClose(await accumulatedGradientsMap[x.id].data(), [0, 5, 0]);
        expectArraysClose(await dys[0].data(), [0]);
        expectArraysClose(await dys[1].data(), [5]);
        expectArraysClose(await dys[2].data(), [0]);
    });
});
//# sourceMappingURL=tape_test.js.map