Fix 1) tanh underflow error, 2) scoping bug and add 3) debug mode for math (#34)

* Fix overflow tanh overflow error, scope bug and add debug mode for math

* simplify debug mode api

* fix typo in trig_gpu_test

* Merge master into logical

* Merge master into logical

* add unit tests for math debug mode

* Merge branch 'logical' of https://github.com/PAIR-code/deeplearnjs into logical

Author: dsmilkov

src/math/math.ts

@@ -28,9 +28,11 @@ export abstract class NDArrayMath {
   private ndarraysToKeep: NDArray[][] = [];
   private activeScopeNDArraysToKeep: NDArray[] = [];
 
+  private debugMode = false;
+
   /**
    * @param safeMode In safe mode, you must use math operations inside
-   * a math.scope() which will automatically clean up intermediate NDArrays.
+   *     a math.scope() which will automatically clean up intermediate NDArrays.
    */
   constructor(private safeMode: boolean) {}
 
@@ -57,6 +59,18 @@ export abstract class NDArrayMath {
     return result;
   }
 
+
+  /**
+   * In debug mode, the output of every math call will be downloaded to the CPU
+   * and checked for NaNs. This significantly impacts performance.
+   */
+  enableDebugMode() {
+    this.debugMode = true;
+    console.warn('Debugging mode is ON. The output of every math call will ' +
+                  'be downloaded to CPU and checked for NaNs. ' +
+                  'This significantly impacts performance.');
+  }
+
   /**
    * Start a scope. Use this with endScope() to achieve the same functionality
    * as scope() without the need for a function closure.
@@ -76,13 +90,14 @@ export abstract class NDArrayMath {
    * as scope() without the need for a function closure.
    */
   endScope(result: ScopeResult) {
+    let arraysToKeep = this.activeScopeNDArraysToKeep;
+    if (result != null) {
+      arraysToKeep = arraysToKeep.concat(result as NDArray|NDArray[]);
+    }
     // Dispose the current scope.
     for (let i = 0; i < this.activeScope.length; i++) {
       const ndarray = this.activeScope[i];
-
-      if (this.isNDArrayDataInList(ndarray, this.activeScopeNDArraysToKeep) ||
-          (result != null && result instanceof NDArray &&
-           ndarray.getData() === (result as NDArray).getData())) {
+      if (this.isNDArrayDataInList(ndarray, arraysToKeep)) {
         continue;
       }
       ndarray.dispose();
@@ -141,12 +156,24 @@ export abstract class NDArrayMath {
     return result;
   }
 
+  private checkForNaN(arr: NDArray): void {
+    const vals = arr.getValues();
+    for (let i = 0; i < vals.length; i++) {
+      if (isNaN(vals[i])) {
+        throw Error('The result NDArray of the last math call has NaNs.');
+      }
+    }
+  }
+
   /**
    * Tracks an NDArray in the current scope to be automatically cleaned up when
    * the current scope ends, and returns the value.
    * @param result The NDArray to track in the current scope.
    */
   track<T extends NDArray>(result: T): T {
+    if (this.debugMode) {
+      this.checkForNaN(result);
+    }
     if (this.activeScope == null) {
       if (this.safeMode) {
         throw new Error(

src/math/math_gpu_test.ts

@@ -27,7 +27,7 @@ describe('NDArrayMathGPU scope', () => {
     math = new NDArrayMathGPU();
   });
 
-  it('basic scope usage with a return', () => {
+  it('scope returns NDArray', () => {
     const a = Array1D.new([1, 2, 3]);
     let b = Array1D.new([0, 0, 0]);
 
@@ -55,6 +55,33 @@ describe('NDArrayMathGPU scope', () => {
     b.dispose();
   });
 
+  it('scope returns NDArray[]', () => {
+    const a = Array1D.new([1, 2, 3]);
+    const b = Array1D.new([0, -1, 1]);
+
+    const numUsedTexturesBefore = math.getTextureManager().getNumUsedTextures();
+
+    math.scope(() => {
+      const result = math.scope(() => {
+        math.add(a, b);
+        return [math.add(a, b), math.sub(a, b)];
+      });
+
+      // a, b, and 2 results are new textures. All intermediates should be
+      // disposed.
+      expect(math.getTextureManager().getNumUsedTextures())
+          .toEqual(numUsedTexturesBefore + 4);
+      expect(result[0].getValues()).toEqual(new Float32Array([1, 1, 4]));
+      expect(result[1].getValues()).toEqual(new Float32Array([1, 3, 2]));
+    });
+
+    // a, b are new textures, result should be disposed.
+    expect(math.getTextureManager().getNumUsedTextures())
+        .toEqual(numUsedTexturesBefore + 2);
+    a.dispose();
+    b.dispose();
+  });
+
   it('basic scope usage without return', () => {
     const a = Array1D.new([1, 2, 3]);
     let b = Array1D.new([0, 0, 0]);
@@ -2128,3 +2155,36 @@ describe('NDArrayMathGPU batchNorm', () => {
     offset.dispose();
   });
 });
+
+describe('NDArrayMathGPU debug mode', () => {
+  let math: NDArrayMathGPU;
+
+  beforeEach(() => {
+    math = new NDArrayMathGPU();
+    math.startScope();
+  });
+
+  afterEach(() => {
+    math.endScope(null!);
+  });
+
+  it('debug mode does not error when no nans', () => {
+    math.enableDebugMode();
+    const a = Array1D.new([2, -1, 0, 3]);
+    const res = math.relu(a);
+    expect(res.getValues()).toEqual(new Float32Array([2, 0, 0, 3]));
+  });
+
+  it('debug mode errors when there are nans', () => {
+    math.enableDebugMode();
+    const a = Array1D.new([2, NaN]);
+    const f = () => math.relu(a);
+    expect(f).toThrowError();
+  });
+
+  it('no errors where there are nans, and debug mode is disabled', () => {
+    const a = Array1D.new([2, NaN]);
+    const res = math.relu(a);
+    expect(res.getValues()).toEqual(new Float32Array([2, NaN]));
+  });
+});

src/math/webgl/trig_gpu_test.ts

@@ -17,7 +17,7 @@ import * as test_util from '../../test_util';
 import * as util from '../../util';
 import {UnaryOp} from './unaryop_gpu';
 import * as unaryop_gpu_test from './unaryop_gpu_test';
-import {Array1D, Array2D, Array3D} from '../ndarray';
+import {Scalar, Array1D, Array2D, Array3D} from '../ndarray';
 
 describe('sin_gpu', () => {
   it('returns a matrix with the same shape as the input matrix', () => {
@@ -58,4 +58,28 @@ describe('tanh_gpu', () => {
     const result = unaryop_gpu_test.uploadUnaryDownload(aArr, UnaryOp.TANH);
     test_util.expectArraysClose(result, expectedResult, 1e-6);
   });
+
+  it('overflow', () => {
+    const a = Scalar.new(100);
+    const r = unaryop_gpu_test.uploadUnaryDownload(a, UnaryOp.TANH);
+    expect(r).toBeCloseTo(1);
+  });
+
+  it('tanh(0) = 0', () => {
+    const a = Scalar.new(0);
+    const r = unaryop_gpu_test.uploadUnaryDownload(a, UnaryOp.TANH);
+    expect(r).toBeCloseTo(0);
+  });
+
+  it('tanh(0.01) is close to 0.01', () => {
+    const a = Scalar.new(0.01);
+    const r = unaryop_gpu_test.uploadUnaryDownload(a, UnaryOp.TANH);
+    expect(r).toBeCloseTo(0.01);
+  });
+
+  it('underflow', () => {
+    const a = Scalar.new(-100);
+    const r = unaryop_gpu_test.uploadUnaryDownload(a, UnaryOp.TANH);
+    expect(r).toBeCloseTo(-1);
+  });
 });

src/math/webgl/unaryop_gpu.ts

@@ -55,8 +55,8 @@ function getOpSnippet(op: UnaryOp) {
     case UnaryOp.SIN:
       return 'float r = sin(v);';
     case UnaryOp.TANH:
-      return `float e2x = exp(-2.0 * v);
-              float r = (1.0 - e2x) / (1.0 + e2x);`;
+      return `float e2x = exp(-2.0 * abs(v));
+              float r = sign(v) * (1.0 - e2x) / (1.0 + e2x);`;
     default:
       throw Error('Unrecognized unary op type ' + op);
   }