Fix bug with dl.customGrad() and add much better error handling for gradient-related API (#731)

Author: dsmilkov

src/engine.ts

@@ -37,13 +37,13 @@ interface ScopeState {
 }
 
 /**
- * @docalias (...inputs: Tensor[]) => {
+ * @docalias (a: Tensor, b: Tensor,...) => {
  *   value: Tensor,
- *   gradFunc: (dy: Tensor) => Tensor[]
+ *   gradFunc: (dy: Tensor) => Tensor|Tensor[]
  * }
  */
 export type CustomGradientFunc<T extends Tensor> = (...args: Tensor[]) => {
-  value: T, gradFunc: (dy: T) => Tensor[];
+  value: T, gradFunc: (dy: T) => Tensor | Tensor[];
 };
 
 export interface TensorManager {
@@ -287,23 +287,27 @@ export class Engine implements TensorManager {
   }
 
   /**
-   * Returns gradients of `f` w.r.t. each of the `xs`. The gradients returned
-   * are of the same length as `xs`, but some might be null if `f` was not
-   * a function of that `x`. It also takes optional dy to multiply the gradient,
-   * which defaults to `1`.
+   * Returns gradients of `f` with respect to each of the `xs`. The gradients
+   * returned are of the same length as `xs`, but some might be null if `f` was
+   * not a function of that `x`. It also takes optional dy to multiply the
+   * gradient, which defaults to `1`.
    */
-  gradients<T extends Tensor>(f: () => T, xs: Tensor[], dy?: T):
-      {value: T, grads: Tensor[]} {
+  gradients<T extends Tensor>(
+      f: () => T, xs: Tensor[], dy?: T,
+      allowNoGradients = false): {value: T, grads: Tensor[]} {
     return tidy('gradients', () => {
       const y = f();
+      util.assert(
+          y instanceof Tensor,
+          'The result y returned by f() must be a tensor.');
       // Filter out the nodes that don't connect x => y.
       const filteredTape =
           tape_util.getFilteredNodesXToY(this.activeTape, xs, y);
-      if (filteredTape.length === 0 && xs.length > 0) {
+      if (!allowNoGradients && filteredTape.length === 0 && xs.length > 0) {
         throw new Error(
-            `Cannot compute gradient: y is not a function of \`x\`s. ` +
-            `Make sure the xs you are computing gradients with respect ` +
-            `to are used inside the gradient function.`);
+            'Cannot compute gradient of y=f(x) with respect to x. Make sure ' +
+            'that the f you passed encloses all operations that lead from x ' +
+            'to y.');
       }
 
       const accumulatedGradientMap: {[tensorId: number]: Tensor} = {};
@@ -319,21 +323,50 @@ export class Engine implements TensorManager {
 
   customGrad<T extends Tensor>(f: CustomGradientFunc<T>):
       (...args: Tensor[]) => T {
-    this.customGradientDepth++;
-
+    util.assert(
+        util.isFunction(f),
+        'The f passed in customGrad(f) must be a function.');
     return (...inputs: Tensor[]): T => {
-      let gradientsFunc: (dy: T) => Tensor[];
+      util.assert(
+          inputs.every(t => t instanceof Tensor),
+          'The args passed in customGrad(f)(x1, x2,...) must all be tensors');
+      this.customGradientDepth++;
+
+      let gradientsFunc: (dy: T) => Tensor | Tensor[];
       const gradientsMode = true;
       const result = tidy(f.name, () => {
         const {value, gradFunc} = f(...inputs);
+        util.assert(
+            value instanceof Tensor,
+            'The function f passed in customGrad(f) must return an object ' +
+                'where `obj.value` is a tensor');
+        util.assert(
+            util.isFunction(gradFunc),
+            'The function f passed in customGrad(f) must return an object ' +
+                'where `obj.gradFunc` is a function.');
         gradientsFunc = gradFunc;
         return value;
       }, gradientsMode);
 
       this.customGradientDepth--;
 
       if (this.shouldRecord()) {
-        this.addTapeNode(inputs, result, gradientsFunc);
+        const gradFunc = (dy: T): Tensor[] => {
+          const res = gradientsFunc(dy);
+          const grads: Tensor[] = Array.isArray(res) ? res : [res];
+          util.assert(
+              grads.length === inputs.length,
+              'The function f passed in customGrad(f) must return an object ' +
+                  'where `obj.gradFunc` is a function that returns the same ' +
+                  'number of tensors as inputs passed to f(...).');
+          util.assert(
+              grads.every(t => t instanceof Tensor),
+              'The function f passed in customGrad(f) must return an object ' +
+                  'where `obj.gradFunc` is a function that returns a list of ' +
+                  'only tensors.');
+          return grads;
+        };
+        this.addTapeNode(inputs, result, gradFunc);
       }
       return result;
     };

src/engine_test.ts

@@ -233,12 +233,30 @@ describeWithFlags('gradients', ALL_ENVS, () => {
     expectArraysClose(result, [.2, .4]);
   });
 
+  it('calling grad(f) twice works', () => {
+    const grad = dl.grad(x => x.square());
+
+    const result = grad(dl.tensor1d([.1, .2]));
+    const result2 = grad(dl.tensor1d([.1, .4]));
+    expectArraysClose(result, [.2, .4]);
+    expectArraysClose(result2, [.2, .8]);
+  });
+
   it('grads(f)', () => {
     const grads = dl.grads(x => x.square());
     const result = grads([dl.tensor1d([.1, .2])]);
     expectArraysClose(result[0], [.2, .4]);
   });
 
+  it('calling grads(f) twice works', () => {
+    const grads = dl.grads(x => x.square());
+
+    const result = grads([dl.tensor1d([.1, .2])]);
+    const result2 = grads([dl.tensor1d([.1, .4])]);
+    expectArraysClose(result[0], [.2, .4]);
+    expectArraysClose(result2[0], [.2, .8]);
+  });
+
   it('works with reshape', () => {
     const a = dl.tensor2d([1, 2, 3, 4], [2, 2]);
     const exponent = dl.tensor1d([2, 2, 2, 2], 'int32');
@@ -390,7 +408,7 @@ describeWithFlags('customGradient', ALL_ENVS, () => {
 
     const customPow = dl.customGrad(a => {
       const value = dl.pow(a, b);
-      const gradFunc = (dy: Tensor) => [dy.mul(dl.scalar(0.1))];
+      const gradFunc = (dy: Tensor) => dy.mul(dl.scalar(0.1));
       return {value, gradFunc};
     });
 
@@ -409,7 +427,7 @@ describeWithFlags('customGradient', ALL_ENVS, () => {
 
     const customPow = dl.customGrad(a => {
       const value = dl.pow(a, b);
-      const gradFunc = (dy: Tensor) => [dy.mul(a)];
+      const gradFunc = (dy: Tensor) => dy.mul(a);
       return {value, gradFunc};
     });
 
@@ -419,6 +437,18 @@ describeWithFlags('customGradient', ALL_ENVS, () => {
     // First order: dy * a. Second order: dy.
     expectArraysClose(dda, dy);
   });
+
+  it('calling gradient of custom op twice works', () => {
+    const customOp = dl.customGrad(x => {
+      // Override gradient of our custom x ^ 2 op to be dy * abs(x);
+      return {value: x.square(), gradFunc: dy => dy.mul(x.abs())};
+    });
+    const x = dl.tensor1d([-1, -2, 3]);
+    const grad = dl.grad(x => customOp(x));
+
+    expectArraysClose(grad(x), [1, 2, 3]);
+    expectArraysClose(grad(x), [1, 2, 3]);
+  });
 });
 
 describeWithFlags('memory', ALL_ENVS, () => {