[core]Modularize addN.

tfjs-core/src/engine.ts

@@ -669,10 +669,22 @@ export class Engine implements TensorTracker, DataMover {
       const inputsToSave: string[] = gradConfig.inputsToSave || [];
       const outputsToSave: boolean[] = gradConfig.outputsToSave || [];
 
-      const inputTensorsToSave: Tensor[] =
-          inputsToSave.map((inputName) => inputs[inputName]);
+      // If saveAllInputs is true, all inputs will be saved. Otherwise, inputs
+      // specified in inputsToSave will be saved.
+      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

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,
+  saveAllInputs: true,
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const ders: {[key: string]: () => Tensor} = {};
+    saved.forEach((_, i) => {
+      ders[i] = () => dy.clone();
+    });
+    return ders;
+  }
+};

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<NamedTensorInfoMap, 'a'|'b'>;
 
 export const Div = 'Div';

tfjs-core/src/kernel_registry.ts

@@ -60,6 +60,9 @@ export interface KernelConfig {
 export interface GradConfig {
   kernelName: string;
   inputsToSave?: string[];
+  // 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;
 }

tfjs-core/src/kernel_registry_test.ts

@@ -258,6 +258,100 @@ 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 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: () => {
+           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);
+           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);
+       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';
     const x = tf.zeros([2, 2]);

tfjs-core/src/ops/add_n.ts

@@ -0,0 +1,77 @@
+/**
+ * @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_<T extends Tensor>(tensors: Array<T|TensorLike>): 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 forward: ForwardFunc<Tensor> = (backend, save) =>
+      backend.addN($tensors);
+
+  const inputs: AddNInputs = $tensors;
+
+  return ENGINE.runKernelFunc(
+             forward, inputs as {} as NamedTensorMap, null /* grad */,
+             'AddN') as T;
+}
+
+export const addN = op({addN_});

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/);
+  });
+});