From a24796a18b5d18fd622bda559ebb93edcf063f93 Mon Sep 17 00:00:00 2001 From: Na Li Date: Tue, 31 Mar 2020 16:47:12 -0700 Subject: [PATCH 1/7] [core]Modularize addN. --- tfjs-core/src/kernel_names.ts | 5 +- tfjs-core/src/ops/add_n.ts | 84 +++++++++++++++++++++++++++ tfjs-core/src/ops/add_n_test.ts | 87 ++++++++++++++++++++++++++++ tfjs-core/src/ops/arithmetic_test.ts | 67 --------------------- tfjs-core/src/ops/binary_ops.ts | 51 ---------------- tfjs-core/src/ops/ops.ts | 1 + tfjs-core/src/tests.ts | 1 + 7 files changed, 177 insertions(+), 119 deletions(-) create mode 100644 tfjs-core/src/ops/add_n.ts create mode 100644 tfjs-core/src/ops/add_n_test.ts diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts index 3d9ce3b440c..20c0b8b3bfb 100644 --- a/tfjs-core/src/kernel_names.ts +++ b/tfjs-core/src/kernel_names.ts @@ -18,12 +18,15 @@ // tslint:disable: variable-name // Unfortunately just enabling PascalCase per file (tslint:enable: // allow-pascal-case) doesn't work. -import {NamedTensorInfoMap} from './kernel_registry'; +import {NamedTensorInfoMap, TensorInfo} from './kernel_registry'; import {PixelData} from './types'; export const Add = 'Add'; export type AddInputs = BinaryInputs; +export const AddN = 'AddN'; +export type AddNInputs = TensorInfo[]; + export type BinaryInputs = Pick; export const Div = 'Div'; diff --git a/tfjs-core/src/ops/add_n.ts b/tfjs-core/src/ops/add_n.ts new file mode 100644 index 00000000000..b348dcf6c19 --- /dev/null +++ b/tfjs-core/src/ops/add_n.ts @@ -0,0 +1,84 @@ +/** + * @license + * Copyright 2020 Google Inc. 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 {ENGINE, ForwardFunc} from '../engine'; +import {AddNInputs} from '../kernel_names'; +import {Tensor} from '../tensor'; +import {NamedTensorMap} from '../tensor_types'; +import {convertToTensor} from '../tensor_util_env'; +import {TensorLike} from '../types'; +import * as util from '../util'; + +import {op} from './operation'; + +/** + * Adds a list of `tf.Tensor`s element-wise, each with the same shape and dtype. + * + * ```js + * const a = tf.tensor1d([1, 2]); + * const b = tf.tensor1d([3, 4]); + * const c = tf.tensor1d([5, 6]); + * + * tf.addN([a, b, c]).print(); + * ``` + * @param tensors A list of tensors with the same shape and dtype. + */ +/** @doc {heading: 'Operations', subheading: 'Arithmetic'} */ +function addN_(tensors: Array): T { + util.assert( + Array.isArray(tensors), + () => 'The argument passed to tf.addN() must be a list of tensors'); + util.assert( + tensors.length >= 1, + () => `Must pass at least one tensor to tf.addN(), but got ` + + `${tensors.length}`); + + const $tensors = + tensors.map((t, i) => convertToTensor(t, `tensors${i}`, 'addN')); + + const firstTensor = $tensors[0]; + $tensors.forEach(t => { + if (t.dtype !== firstTensor.dtype) { + throw new Error( + 'All tensors passed to tf.addN() must have the same dtype'); + } + }); + + $tensors.forEach(t => { + if (!util.arraysEqual(t.shape, firstTensor.shape)) { + throw new Error( + 'All tensors passed to tf.addN() must have the same shape'); + } + }); + + const der = (dy: T) => { + const ders: {[key: string]: () => Tensor} = {}; + $tensors.forEach((t, i) => { + ders[i] = () => dy.clone(); + }); + return ders; + }; + + const forward: ForwardFunc = (backend, save) => + backend.addN($tensors); + + const inputs: AddNInputs = $tensors; + + return ENGINE.runKernelFunc( + forward, inputs as {} as NamedTensorMap, der, 'AddN') as T; +} + +export const addN = op({addN_}); diff --git a/tfjs-core/src/ops/add_n_test.ts b/tfjs-core/src/ops/add_n_test.ts new file mode 100644 index 00000000000..a999286406f --- /dev/null +++ b/tfjs-core/src/ops/add_n_test.ts @@ -0,0 +1,87 @@ +/** + * @license + * Copyright 2020 Google Inc. 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 * as tf from '../index'; +import {ALL_ENVS, describeWithFlags} from '../jasmine_util'; +import {expectArraysClose} from '../test_util'; + +describeWithFlags('addN', ALL_ENVS, () => { + it('a single tensor', async () => { + const res = tf.addN([tf.tensor1d([1, 2, 3])]); + expectArraysClose(await res.data(), [1, 2, 3]); + }); + + it('two tensors, int32', async () => { + const res = tf.addN([ + tf.tensor1d([1, 2, -1], 'int32'), + tf.tensor1d([5, 3, 2], 'int32'), + ]); + expectArraysClose(await res.data(), [6, 5, 1]); + expect(res.dtype).toBe('int32'); + expect(res.shape).toEqual([3]); + }); + + it('three tensors', async () => { + const res = tf.addN([ + tf.tensor1d([1, 2]), + tf.tensor1d([5, 3]), + tf.tensor1d([-5, -2]), + ]); + expectArraysClose(await res.data(), [1, 3]); + expect(res.dtype).toBe('float32'); + expect(res.shape).toEqual([2]); + }); + + it('accepts a tensor-like object', async () => { + const res = tf.addN([[1, 2], [3, 4]]); + expectArraysClose(await res.data(), [4, 6]); + expect(res.dtype).toBe('float32'); + expect(res.shape).toEqual([2]); + }); + + it('list of numbers gets treated as a list of scalars', async () => { + const res = tf.addN([1, 2, 3, 4]); + expectArraysClose(await res.data(), [10]); + expect(res.dtype).toBe('float32'); + expect(res.shape).toEqual([]); + }); + + it('errors if list is empty', () => { + expect(() => tf.addN([])) + .toThrowError( + /Must pass at least one tensor to tf.addN\(\), but got 0/); + }); + + it('errors if argument is not an array', () => { + // tslint:disable-next-line:no-any + expect(() => tf.addN(tf.scalar(3) as any)) + .toThrowError( + /The argument passed to tf.addN\(\) must be a list of tensors/); + }); + + it('errors if arguments not of same dtype', () => { + expect(() => tf.addN([tf.scalar(1, 'int32'), tf.scalar(2, 'float32')])) + .toThrowError( + /All tensors passed to tf.addN\(\) must have the same dtype/); + }); + + it('errors if arguments not of same shape', () => { + expect(() => tf.addN([tf.scalar(1), tf.tensor1d([2])])) + .toThrowError( + /All tensors passed to tf.addN\(\) must have the same shape/); + }); +}); diff --git a/tfjs-core/src/ops/arithmetic_test.ts b/tfjs-core/src/ops/arithmetic_test.ts index aeda0b01dbf..12c7e38152a 100644 --- a/tfjs-core/src/ops/arithmetic_test.ts +++ b/tfjs-core/src/ops/arithmetic_test.ts @@ -1014,73 +1014,6 @@ describeWithFlags('pow', ALL_ENVS, () => { }); }); -describeWithFlags('addN', ALL_ENVS, () => { - it('a single tensor', async () => { - const res = tf.addN([tf.tensor1d([1, 2, 3])]); - expectArraysClose(await res.data(), [1, 2, 3]); - }); - - it('two tensors, int32', async () => { - const res = tf.addN([ - tf.tensor1d([1, 2, -1], 'int32'), - tf.tensor1d([5, 3, 2], 'int32'), - ]); - expectArraysClose(await res.data(), [6, 5, 1]); - expect(res.dtype).toBe('int32'); - expect(res.shape).toEqual([3]); - }); - - it('three tensors', async () => { - const res = tf.addN([ - tf.tensor1d([1, 2]), - tf.tensor1d([5, 3]), - tf.tensor1d([-5, -2]), - ]); - expectArraysClose(await res.data(), [1, 3]); - expect(res.dtype).toBe('float32'); - expect(res.shape).toEqual([2]); - }); - - it('accepts a tensor-like object', async () => { - const res = tf.addN([[1, 2], [3, 4]]); - expectArraysClose(await res.data(), [4, 6]); - expect(res.dtype).toBe('float32'); - expect(res.shape).toEqual([2]); - }); - - it('list of numbers gets treated as a list of scalars', async () => { - const res = tf.addN([1, 2, 3, 4]); - expectArraysClose(await res.data(), [10]); - expect(res.dtype).toBe('float32'); - expect(res.shape).toEqual([]); - }); - - it('errors if list is empty', () => { - expect(() => tf.addN([])) - .toThrowError( - /Must pass at least one tensor to tf.addN\(\), but got 0/); - }); - - it('errors if argument is not an array', () => { - // tslint:disable-next-line:no-any - expect(() => tf.addN(tf.scalar(3) as any)) - .toThrowError( - /The argument passed to tf.addN\(\) must be a list of tensors/); - }); - - it('errors if arguments not of same dtype', () => { - expect(() => tf.addN([tf.scalar(1, 'int32'), tf.scalar(2, 'float32')])) - .toThrowError( - /All tensors passed to tf.addN\(\) must have the same dtype/); - }); - - it('errors if arguments not of same shape', () => { - expect(() => tf.addN([tf.scalar(1), tf.tensor1d([2])])) - .toThrowError( - /All tensors passed to tf.addN\(\) must have the same shape/); - }); -}); - describeWithFlags('sub', ALL_ENVS, () => { it('c - A', async () => { const c = tf.scalar(5); diff --git a/tfjs-core/src/ops/binary_ops.ts b/tfjs-core/src/ops/binary_ops.ts index c44c82af29c..6699fe48557 100644 --- a/tfjs-core/src/ops/binary_ops.ts +++ b/tfjs-core/src/ops/binary_ops.ts @@ -17,7 +17,6 @@ import {ENGINE} from '../engine'; import {Tensor} from '../tensor'; -import {NamedTensorMap} from '../tensor_types'; import {makeTypesMatch} from '../tensor_util'; import {convertToTensor} from '../tensor_util_env'; import {TensorLike} from '../types'; @@ -29,55 +28,6 @@ import {op} from './operation'; import {scalar, zerosLike} from './tensor_ops'; import {neg} from './unary_ops'; -/** - * Adds a list of `tf.Tensor`s element-wise, each with the same shape and dtype. - * - * ```js - * const a = tf.tensor1d([1, 2]); - * const b = tf.tensor1d([3, 4]); - * const c = tf.tensor1d([5, 6]); - * - * tf.addN([a, b, c]).print(); - * ``` - * @param tensors A list of tensors with the same shape and dtype. - */ -/** @doc {heading: 'Operations', subheading: 'Arithmetic'} */ -function addN_(tensors: Array): T { - util.assert( - Array.isArray(tensors), - () => 'The argument passed to tf.addN() must be a list of tensors'); - util.assert( - tensors.length >= 1, - () => `Must pass at least one tensor to tf.addN(), but got ` + - `${tensors.length}`); - const $tensors = - tensors.map((t, i) => convertToTensor(t, `tensors${i}`, 'addN')); - const firstTensor = $tensors[0]; - $tensors.forEach(t => { - if (t.dtype !== firstTensor.dtype) { - throw new Error( - 'All tensors passed to tf.addN() must have the same dtype'); - } - }); - $tensors.forEach(t => { - if (!util.arraysEqual(t.shape, firstTensor.shape)) { - throw new Error( - 'All tensors passed to tf.addN() must have the same shape'); - } - }); - - const der = (dy: T) => { - const ders: {[key: string]: () => Tensor} = {}; - $tensors.forEach((t, i) => { - ders[i] = () => dy.clone(); - }); - return ders; - }; - const inputs: NamedTensorMap = $tensors as {} as NamedTensorMap; - return ENGINE.runKernelFunc( - backend => backend.addN($tensors), inputs, der, 'AddN'); -} - /** * Adds two `tf.Tensor`s element-wise, A + B. * @@ -675,7 +625,6 @@ function atan2_( }, {$a, $b}, der) as T; } -export const addN = op({addN_}); export const addStrict = op({addStrict_}); export const atan2 = op({atan2_}); export const divStrict = op({divStrict_}); diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts index 32dcfac224d..3b5f0f59c18 100644 --- a/tfjs-core/src/ops/ops.ts +++ b/tfjs-core/src/ops/ops.ts @@ -17,6 +17,7 @@ // Modularized ops. export {add} from './add'; +export {addN} from './add_n'; export {batchNorm, batchNormalization} from './batchnorm'; export {batchNorm2d, batchNormalization2d} from './batchnorm2d'; export {batchNorm3d, batchNormalization3d} from './batchnorm3d'; diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index a6515ae1622..6aaf209522e 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -38,6 +38,7 @@ import './io/weights_loader_test'; import './jasmine_util_test'; import './kernel_registry_test'; import './ops/add_test'; +import './ops/add_n_test'; import './ops/arithmetic_test'; import './ops/array_ops_test'; import './ops/axis_util_test'; From a179accf7070e899a100418d84ae59c1f5acd17c Mon Sep 17 00:00:00 2001 From: Na Li Date: Wed, 1 Apr 2020 17:48:57 -0700 Subject: [PATCH 2/7] Separate out gradient. --- tfjs-core/src/engine.ts | 15 +++++++++--- tfjs-core/src/gradients/AddN_grad.ts | 32 +++++++++++++++++++++++++ tfjs-core/src/kernel_registry.ts | 3 +++ tfjs-core/src/ops/add_n.ts | 11 ++------- tfjs-core/src/register_all_gradients.ts | 8 ++++--- 5 files changed, 54 insertions(+), 15 deletions(-) create mode 100644 tfjs-core/src/gradients/AddN_grad.ts diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 3c518d9be37..bd93d749e73 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -666,11 +666,20 @@ export class Engine implements TensorTracker, DataMover { outputs: Tensor[]): Tensor[]|null { const gradConfig = getGradient(kernelName); if (gradConfig != null) { - const inputsToSave: string[] = gradConfig.inputsToSave || []; + const inputsToSave: string[] = gradConfig.inputsToSave; const outputsToSave: boolean[] = gradConfig.outputsToSave || []; - const inputTensorsToSave: Tensor[] = - inputsToSave.map((inputName) => inputs[inputName]); + let inputTensorsToSave: Tensor[]; + if (!inputsToSave) { + inputTensorsToSave = []; + } else if (inputsToSave.length === 0) { + // An empty inputsToSave array will result in save all the inputs. + inputTensorsToSave = Object.keys(inputs).map((key) => inputs[key]); + } else { + // A non-empty inputsToSave array will only save specified inputs. + inputTensorsToSave = inputsToSave.map((inputName) => inputs[inputName]); + } + const outputTensorsToSave: Tensor[] = outputs.filter((_, i) => outputsToSave[i]); return inputTensorsToSave.concat(outputTensorsToSave); diff --git a/tfjs-core/src/gradients/AddN_grad.ts b/tfjs-core/src/gradients/AddN_grad.ts new file mode 100644 index 00000000000..6b08fcdf3af --- /dev/null +++ b/tfjs-core/src/gradients/AddN_grad.ts @@ -0,0 +1,32 @@ +/** + * @license + * Copyright 2020 Google Inc. 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 {AddN} from '../kernel_names'; +import {GradConfig} from '../kernel_registry'; +import {Tensor} from '../tensor'; + +export const addNGradConfig: GradConfig = { + kernelName: AddN, + inputsToSave: [], + gradFunc: (dy: Tensor, saved: Tensor[]) => { + const ders: {[key: string]: () => Tensor} = {}; + saved.forEach((_, i) => { + ders[i] = () => dy.clone(); + }); + return ders; + } +}; diff --git a/tfjs-core/src/kernel_registry.ts b/tfjs-core/src/kernel_registry.ts index 3997636ca3d..f2124f09f42 100644 --- a/tfjs-core/src/kernel_registry.ts +++ b/tfjs-core/src/kernel_registry.ts @@ -59,6 +59,9 @@ export interface KernelConfig { /** Config object for registering a gradient in the global registry. */ export interface GradConfig { kernelName: string; + // If the array is empty, all the inputs will be saved. + // If the array is not empty, only the specified inputs will be saved. + // If the array is undefined, no inputs will be saved. inputsToSave?: string[]; outputsToSave?: boolean[]; gradFunc: GradFunc; diff --git a/tfjs-core/src/ops/add_n.ts b/tfjs-core/src/ops/add_n.ts index b348dcf6c19..994c7ccd17a 100644 --- a/tfjs-core/src/ops/add_n.ts +++ b/tfjs-core/src/ops/add_n.ts @@ -64,21 +64,14 @@ function addN_(tensors: Array): T { } }); - const der = (dy: T) => { - const ders: {[key: string]: () => Tensor} = {}; - $tensors.forEach((t, i) => { - ders[i] = () => dy.clone(); - }); - return ders; - }; - const forward: ForwardFunc = (backend, save) => backend.addN($tensors); const inputs: AddNInputs = $tensors; return ENGINE.runKernelFunc( - forward, inputs as {} as NamedTensorMap, der, 'AddN') as T; + forward, inputs as {} as NamedTensorMap, null /* grad */, + 'AddN') as T; } export const addN = op({addN_}); diff --git a/tfjs-core/src/register_all_gradients.ts b/tfjs-core/src/register_all_gradients.ts index 62d0c1b738a..3ce07d7d446 100644 --- a/tfjs-core/src/register_all_gradients.ts +++ b/tfjs-core/src/register_all_gradients.ts @@ -15,6 +15,7 @@ * ============================================================================= */ import {addGradConfig} from './gradients/Add_grad'; +import {addNGradConfig} from './gradients/AddN_grad'; import {broadcastToGradConfig} from './gradients/BroadcastTo_grad'; import {divGradConfig} from './gradients/Div_grad'; import {fusedBatchNormGradConfig} from './gradients/FusedBatchNorm_grad'; @@ -30,9 +31,10 @@ import {registerGradient} from './kernel_registry'; // Export all kernel configs here so that the package can auto register them const gradConfigs: GradConfig[] = [ - addGradConfig, broadcastToGradConfig, divGradConfig, fusedBatchNormGradConfig, - identityGradConfig, oneHotGradConfig, padV2GradConfig, squareGradConfig, - squaredDifferenceGradConfig, tileGradConfig, transposeGradConfig + addGradConfig, addNGradConfig, broadcastToGradConfig, divGradConfig, + fusedBatchNormGradConfig, identityGradConfig, oneHotGradConfig, + padV2GradConfig, squareGradConfig, squaredDifferenceGradConfig, + tileGradConfig, transposeGradConfig ]; for (const gradientConfig of gradConfigs) { From b30595a6d1683212662c51c66941d0b211b17349 Mon Sep 17 00:00:00 2001 From: Na Li Date: Wed, 1 Apr 2020 17:55:57 -0700 Subject: [PATCH 3/7] . --- tfjs-core/src/engine.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index bd93d749e73..cd5d71f380d 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -671,12 +671,13 @@ export class Engine implements TensorTracker, DataMover { let inputTensorsToSave: Tensor[]; if (!inputsToSave) { + // If the array is undefined, no inputs will be saved. inputTensorsToSave = []; } else if (inputsToSave.length === 0) { - // An empty inputsToSave array will result in save all the inputs. + // If the array is empty, all the inputs will be saved. inputTensorsToSave = Object.keys(inputs).map((key) => inputs[key]); } else { - // A non-empty inputsToSave array will only save specified inputs. + // If the array is not empty, only the specified inputs will be saved. inputTensorsToSave = inputsToSave.map((inputName) => inputs[inputName]); } From cf427950d6dd903e4235e92dc34c6d6dde82eb60 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 2 Apr 2020 11:41:01 -0700 Subject: [PATCH 4/7] Add a boolean flag to gradConfig to save all inputs to handle inputs with dynamic number of properties. --- scripts/release-notes.ts | 0 tfjs-core/src/engine.ts | 19 ++++++------------- tfjs-core/src/gradients/AddN_grad.ts | 2 +- tfjs-core/src/kernel_registry.ts | 6 +++--- 4 files changed, 10 insertions(+), 17 deletions(-) create mode 100644 scripts/release-notes.ts diff --git a/scripts/release-notes.ts b/scripts/release-notes.ts new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index cd5d71f380d..7093f6436b7 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -666,21 +666,14 @@ export class Engine implements TensorTracker, DataMover { outputs: Tensor[]): Tensor[]|null { const gradConfig = getGradient(kernelName); if (gradConfig != null) { - const inputsToSave: string[] = gradConfig.inputsToSave; + const inputsToSave: string[] = gradConfig.inputsToSave || []; const outputsToSave: boolean[] = gradConfig.outputsToSave || []; - let inputTensorsToSave: Tensor[]; - if (!inputsToSave) { - // If the array is undefined, no inputs will be saved. - inputTensorsToSave = []; - } else if (inputsToSave.length === 0) { - // If the array is empty, all the inputs will be saved. - inputTensorsToSave = Object.keys(inputs).map((key) => inputs[key]); - } else { - // If the array is not empty, only the specified inputs will be saved. - inputTensorsToSave = inputsToSave.map((inputName) => inputs[inputName]); - } - + // If saveAllInputs is true, all inputs will be saved. Otherwise, inputs + // specified in inputsToSave will be saved. + const inputTensorsToSave: Tensor[] = gradConfig.saveAllInputs ? + Object.keys(inputs).map((key) => inputs[key]) : + inputsToSave.map((inputName) => inputs[inputName]); const outputTensorsToSave: Tensor[] = outputs.filter((_, i) => outputsToSave[i]); return inputTensorsToSave.concat(outputTensorsToSave); diff --git a/tfjs-core/src/gradients/AddN_grad.ts b/tfjs-core/src/gradients/AddN_grad.ts index 6b08fcdf3af..77fd248588e 100644 --- a/tfjs-core/src/gradients/AddN_grad.ts +++ b/tfjs-core/src/gradients/AddN_grad.ts @@ -21,7 +21,7 @@ import {Tensor} from '../tensor'; export const addNGradConfig: GradConfig = { kernelName: AddN, - inputsToSave: [], + saveAllInputs: true, gradFunc: (dy: Tensor, saved: Tensor[]) => { const ders: {[key: string]: () => Tensor} = {}; saved.forEach((_, i) => { diff --git a/tfjs-core/src/kernel_registry.ts b/tfjs-core/src/kernel_registry.ts index f2124f09f42..10f23d0457b 100644 --- a/tfjs-core/src/kernel_registry.ts +++ b/tfjs-core/src/kernel_registry.ts @@ -59,10 +59,10 @@ export interface KernelConfig { /** Config object for registering a gradient in the global registry. */ export interface GradConfig { kernelName: string; - // If the array is empty, all the inputs will be saved. - // If the array is not empty, only the specified inputs will be saved. - // If the array is undefined, no inputs will be saved. inputsToSave?: string[]; + // If saveAllInputs is true, inputsToSave will be ignored, all inputs will + // be saved. + saveAllInputs?: boolean; outputsToSave?: boolean[]; gradFunc: GradFunc; } From f5db7da6daf398efcb07821cf89bd2781e770166 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 2 Apr 2020 15:19:15 -0700 Subject: [PATCH 5/7] Add tests. --- scripts/release-notes.ts | 0 tfjs-core/src/kernel_registry_test.ts | 99 +++++++++++++++++++++++++++ 2 files changed, 99 insertions(+) delete mode 100644 scripts/release-notes.ts diff --git a/scripts/release-notes.ts b/scripts/release-notes.ts deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tfjs-core/src/kernel_registry_test.ts b/tfjs-core/src/kernel_registry_test.ts index b7afde14e8e..7275b24ed4b 100644 --- a/tfjs-core/src/kernel_registry_test.ts +++ b/tfjs-core/src/kernel_registry_test.ts @@ -258,6 +258,105 @@ describeWithFlags('gradient registry', ALL_ENVS, () => { tf.unregisterGradient(kernelName); }); + it('register a kernel with array inputs and saveAllInputs true', async () => { + let kernelWasCalled = false; + let gradientWasCalled = false; + const kernelName = 'MyKernel'; + const x = [tf.zeros([2, 2]), tf.zeros([2, 2])]; + + const forwardReturnDataId = {}; + tf.registerKernel({ + kernelName, + backendName: tf.getBackend(), + kernelFunc: () => { + kernelWasCalled = true; + return {dtype: 'float32', shape: [3, 3], dataId: forwardReturnDataId}; + } + }); + + tf.registerGradient({ + kernelName, + saveAllInputs: true, + gradFunc: (dy: tf.Tensor, saved) => { + // Make sure saved input (x) was passed to the gradient function. + const [$x0, $x1] = x; + expect(saved.length).toEqual(x.length); + expect($x0.dataId).toEqual(x[0].dataId); + expect($x1.dataId).toEqual(x[1].dataId); + gradientWasCalled = true; + return {0: () => tf.fill([2, 2], 3), 1: () => tf.fill([2, 2], 3)}; + } + }); + + // Inputs as array. + const z = (...x: tf.Tensor[]) => + tf.engine().runKernel( + kernelName, x as {} as tf.NamedTensorMap, {} /* attrs */) as + tf.Tensor; + const gradFunc = tf.grads(z); + const dx = gradFunc(x); + expect(kernelWasCalled).toBe(true); + expect(gradientWasCalled).toBe(true); + expect(dx.length).toEqual(2); + expect(dx[0].dtype).toBe('float32'); + expect(dx[0].shape).toEqual([2, 2]); + expect(dx[1].dtype).toBe('float32'); + expect(dx[1].shape).toEqual([2, 2]); + expectArraysClose(await dx[0].data(), [3, 3, 3, 3]); + expectArraysClose(await dx[1].data(), [3, 3, 3, 3]); + tf.unregisterKernel(kernelName, tf.getBackend()); + tf.unregisterGradient(kernelName); + }); + + it('register a kernel with map inputs and saveAllInputs true', async () => { + let kernelWasCalled = false; + let gradientWasCalled = false; + const kernelName = 'MyKernel'; + const x0 = tf.zeros([2, 2]); + const x1 = tf.zeros([2, 2]); + + const forwardReturnDataId = {}; + tf.registerKernel({ + kernelName, + backendName: tf.getBackend(), + kernelFunc: () => { + kernelWasCalled = true; + return {dtype: 'float32', shape: [3, 3], dataId: forwardReturnDataId}; + } + }); + + tf.registerGradient({ + kernelName, + saveAllInputs: true, + gradFunc: (dy: tf.Tensor, saved) => { + // Make sure saved input (x) was passed to the gradient function. + const [$x0, $x1] = saved; + expect($x0.dataId).toEqual(x0.dataId); + expect($x1.dataId).toEqual(x1.dataId); + gradientWasCalled = true; + return {x0: () => tf.fill([2, 2], 3), x1: () => tf.fill([2, 2], 3)}; + } + }); + + // Inputs as map. + const z = (x0: tf.Tensor, x1: tf.Tensor) => + tf.engine().runKernel(kernelName, {x0, x1}, {} /* attrs */) as + tf.Tensor; + const gradFunc = tf.grads(z); + const dx = gradFunc([x0, x1]); + expect(kernelWasCalled).toBe(true); + expect(gradientWasCalled).toBe(true); + expect(dx.length).toEqual(2); + expect(dx[0].dtype).toBe('float32'); + expect(dx[0].shape).toEqual([2, 2]); + expect(dx[1].dtype).toBe('float32'); + expect(dx[1].shape).toEqual([2, 2]); + // expectArraysClose(await dx[0].data(), [3, 3, 3, 3]); + // expectArraysClose(await dx[1].data(), [3, 3, 3, 3]); + tf.unregisterKernel(kernelName, tf.getBackend()); + tf.unregisterGradient(kernelName); + }); + it('errors when running non-existent gradient', () => { const kernelName = 'MyKernel'; const x = tf.zeros([2, 2]); From 8f612bcd119bb39bb3047aaf22e8a3df82ff0f2e Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 2 Apr 2020 22:52:37 -0700 Subject: [PATCH 6/7] Add more restriction to the flag. --- tfjs-core/src/engine.ts | 15 ++++- tfjs-core/src/kernel_registry.ts | 4 +- tfjs-core/src/kernel_registry_test.ts | 85 +++++++++++++-------------- 3 files changed, 54 insertions(+), 50 deletions(-) diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts index 7093f6436b7..a58ce59e5ca 100644 --- a/tfjs-core/src/engine.ts +++ b/tfjs-core/src/engine.ts @@ -671,11 +671,20 @@ export class Engine implements TensorTracker, DataMover { // If saveAllInputs is true, all inputs will be saved. Otherwise, inputs // specified in inputsToSave will be saved. - const inputTensorsToSave: Tensor[] = gradConfig.saveAllInputs ? - Object.keys(inputs).map((key) => inputs[key]) : - inputsToSave.map((inputName) => inputs[inputName]); + let inputTensorsToSave: Tensor[]; + if (gradConfig.saveAllInputs) { + util.assert( + Array.isArray(inputs), + () => 'saveAllInputs is true, expected inputs to be an array.'); + + inputTensorsToSave = Object.keys(inputs).map((key) => inputs[key]); + } else { + inputTensorsToSave = inputsToSave.map((inputName) => inputs[inputName]); + } + const outputTensorsToSave: Tensor[] = outputs.filter((_, i) => outputsToSave[i]); + return inputTensorsToSave.concat(outputTensorsToSave); } // TODO(yassogba) throw exception here once all runkernelFunc calls with diff --git a/tfjs-core/src/kernel_registry.ts b/tfjs-core/src/kernel_registry.ts index 10f23d0457b..ec946510da5 100644 --- a/tfjs-core/src/kernel_registry.ts +++ b/tfjs-core/src/kernel_registry.ts @@ -60,8 +60,8 @@ export interface KernelConfig { export interface GradConfig { kernelName: string; inputsToSave?: string[]; - // If saveAllInputs is true, inputsToSave will be ignored, all inputs will - // be saved. + // When saveAllInputs is true, all inputs will be saved. Only use this flag + // if inputs is an array of Tensors. saveAllInputs?: boolean; outputsToSave?: boolean[]; gradFunc: GradFunc; diff --git a/tfjs-core/src/kernel_registry_test.ts b/tfjs-core/src/kernel_registry_test.ts index 7275b24ed4b..8478db22eb6 100644 --- a/tfjs-core/src/kernel_registry_test.ts +++ b/tfjs-core/src/kernel_registry_test.ts @@ -308,54 +308,49 @@ describeWithFlags('gradient registry', ALL_ENVS, () => { tf.unregisterGradient(kernelName); }); - it('register a kernel with map inputs and saveAllInputs true', async () => { - let kernelWasCalled = false; - let gradientWasCalled = false; - const kernelName = 'MyKernel'; - const x0 = tf.zeros([2, 2]); - const x1 = tf.zeros([2, 2]); + it('register a kernel with map inputs and saveAllInputs true should throw ' + + 'error', + async () => { + const kernelName = 'MyKernel'; + const x0 = tf.zeros([2, 2]); + const x1 = tf.zeros([2, 2]); - const forwardReturnDataId = {}; - tf.registerKernel({ - kernelName, - backendName: tf.getBackend(), - kernelFunc: () => { - kernelWasCalled = true; - return {dtype: 'float32', shape: [3, 3], dataId: forwardReturnDataId}; - } - }); + const forwardReturnDataId = {}; + tf.registerKernel({ + kernelName, + backendName: tf.getBackend(), + kernelFunc: () => { + return { + dtype: 'float32', + shape: [3, 3], + dataId: forwardReturnDataId + }; + } + }); - tf.registerGradient({ - kernelName, - saveAllInputs: true, - gradFunc: (dy: tf.Tensor, saved) => { - // Make sure saved input (x) was passed to the gradient function. - const [$x0, $x1] = saved; - expect($x0.dataId).toEqual(x0.dataId); - expect($x1.dataId).toEqual(x1.dataId); - gradientWasCalled = true; - return {x0: () => tf.fill([2, 2], 3), x1: () => tf.fill([2, 2], 3)}; - } - }); + tf.registerGradient({ + kernelName, + saveAllInputs: true, + gradFunc: (dy: tf.Tensor, saved) => { + // Make sure saved input (x) was passed to the gradient function. + const [$x0, $x1] = saved; + expect($x0.dataId).toEqual(x0.dataId); + expect($x1.dataId).toEqual(x1.dataId); + return {x0: () => tf.fill([2, 2], 3), x1: () => tf.fill([2, 2], 3)}; + } + }); - // Inputs as map. - const z = (x0: tf.Tensor, x1: tf.Tensor) => - tf.engine().runKernel(kernelName, {x0, x1}, {} /* attrs */) as - tf.Tensor; - const gradFunc = tf.grads(z); - const dx = gradFunc([x0, x1]); - expect(kernelWasCalled).toBe(true); - expect(gradientWasCalled).toBe(true); - expect(dx.length).toEqual(2); - expect(dx[0].dtype).toBe('float32'); - expect(dx[0].shape).toEqual([2, 2]); - expect(dx[1].dtype).toBe('float32'); - expect(dx[1].shape).toEqual([2, 2]); - // expectArraysClose(await dx[0].data(), [3, 3, 3, 3]); - // expectArraysClose(await dx[1].data(), [3, 3, 3, 3]); - tf.unregisterKernel(kernelName, tf.getBackend()); - tf.unregisterGradient(kernelName); - }); + // Inputs as map. + const z = (x0: tf.Tensor, x1: tf.Tensor) => + tf.engine().runKernel(kernelName, {x0, x1}, {} /* attrs */) as + tf.Tensor; + const gradFunc = tf.grads(z); + expect(() => gradFunc([x0, x1])) + .toThrowError( + /saveAllInputs is true, expected inputs to be an array/); + tf.unregisterKernel(kernelName, tf.getBackend()); + tf.unregisterGradient(kernelName); + }); it('errors when running non-existent gradient', () => { const kernelName = 'MyKernel'; From d4fd3c9434d0ec1bd2d2b23814d45adaebdbe2d0 Mon Sep 17 00:00:00 2001 From: Na Li Date: Thu, 2 Apr 2020 23:12:50 -0700 Subject: [PATCH 7/7] . --- tfjs-core/src/tests.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tfjs-core/src/tests.ts b/tfjs-core/src/tests.ts index 6aaf209522e..53ba1e75a5a 100644 --- a/tfjs-core/src/tests.ts +++ b/tfjs-core/src/tests.ts @@ -37,8 +37,8 @@ import './io/router_registry_test'; import './io/weights_loader_test'; import './jasmine_util_test'; import './kernel_registry_test'; -import './ops/add_test'; import './ops/add_n_test'; +import './ops/add_test'; import './ops/arithmetic_test'; import './ops/array_ops_test'; import './ops/axis_util_test';