Add math.multinomial and math.oneHot (#160)

* add multinomial

* add cpu implementation and tests

* add onehot operation

* resolve comments

Author: dsmilkov

package.json

@@ -13,6 +13,7 @@
   "devDependencies": {
     "@types/jasmine": "~2.5.53",
     "@types/polymer": "~1.1.31",
+    "@types/seedrandom": "~2.4.27",
     "bower": "~1.8.0",
     "browserify": "~14.4.0",
     "cross-spawn": "~5.1.0",
@@ -37,5 +38,8 @@
     "build": "tsc",
     "test": "karma start",
     "lint": "tslint -p . --type-check -t verbose"
+  },
+  "dependencies": {
+    "seedrandom": "~2.4.3"
   }
 }

src/math/math.ts

@@ -194,6 +194,9 @@ export abstract class NDArrayMath {
     return result;
   }
 
+  /** Disposes the math object and any resources used by it. */
+  dispose() {}
+
   /**
    * Computes the dot product of two matrices, A * B. These must be matrices,
    * use matrixTimesVector and vectorTimesMatrix, dotProduct, and outerProduct
@@ -270,9 +273,8 @@ export abstract class NDArrayMath {
             `rank ${matrix.rank}.`);
     util.assert(
         v.size === matrix.shape[0],
-        `Error in vectorTimesMatrix: size of first rank 1 input (${v.size}) ` +
-            `must match inner dimension of second rank 2 input, but got ` +
-            `rank ${matrix.rank}.`);
+        `Error in vectorTimesMatrix: size of vector (${v.size}) ` +
+            `must match first dimension of matrix (${matrix.shape[0]})`);
 
     return this.matMul(v.as2D(1, -1), matrix).as1D();
   }
@@ -1574,6 +1576,52 @@ export abstract class NDArrayMath {
     });
     return [res[0].as2D(1, -1), res[1].as2D(1, -1)];
   }
+
+  /**
+   * Draws samples from a multinomial distribution.
+   *
+   * @param probabilities 1D array with normalized outcome probabilities.
+   * @param numSamples Number of samples to draw.
+   * @param seed Optional. The seed number.
+   */
+  multinomial(probabilities: Array1D, numSamples: number, seed?: number):
+      Array1D {
+    const numOutcomes = probabilities.size;
+    if (numOutcomes < 2) {
+      throw new Error(
+          `Error in multinomial: you need at least 2 outcomes, but got ` +
+          `${numOutcomes}.`);
+    }
+    seed = seed || Math.random();
+    return this.executeOp(
+        'multinomial',
+        () => this.multinomialInternal(probabilities, numSamples, seed));
+  }
+  protected abstract multinomialInternal(
+      probabilities: Array1D, numSamples: number, seed: number): Array1D;
+
+  /**
+   * Returns a one-hot tensor. The locations represented by `indices` take
+   * value `onValue` (defaults to 1), while all other locations take value
+   * `offValue` (defaults to 0).
+   *
+   * @param indices 1D Array of indices.
+   * @param depth The depth of the one hot dimension.
+   * @param onValue A number used to fill in output when the index matches the
+   *     location.
+   * @param offValue A number used to fill in the output when the index does not
+   *     match the location.
+   */
+  oneHot(indices: Array1D, depth: number, onValue = 1, offValue = 0): Array2D {
+    if (depth < 2) {
+      throw new Error(`Error in oneHot: depth must be >=2, but it is ${depth}`);
+    }
+    return this.executeOp(
+        'oneHot', () => this.oneHotInternal(indices, depth, onValue, offValue));
+  }
+  protected abstract oneHotInternal(
+      indices: Array1D, depth: number, onValue: number,
+      offValue: number): Array2D;
 }
 
 export enum MatrixOrientation {

src/math/math_cpu.ts

@@ -15,8 +15,8 @@
  * =============================================================================
  */
 
+import * as seedrandom from 'seedrandom';
 import * as util from '../util';
-
 import * as concat_util from './concat_util';
 import * as conv_util from './conv_util';
 import {ConvInfo} from './conv_util';
@@ -31,7 +31,8 @@ export class NDArrayMathCPU extends NDArrayMath {
 
   protected cloneInternal<T extends NDArray>(ndarray: T): T {
     return NDArray.make(
-        ndarray.shape, {values: new Float32Array(ndarray.getValues())}) as T;
+               ndarray.shape,
+               {values: new Float32Array(ndarray.getValues())}) as T;
   }
 
   protected slice1DInternal(input: Array1D, begin: number, size: number):
@@ -982,4 +983,48 @@ export class NDArrayMathCPU extends NDArrayMath {
     }
     return Array3D.make(x.shape, {values: outValues});
   }
+
+  protected multinomialInternal(
+      probabilities: Array1D, numSamples: number, seed: number): Array1D {
+    const probVals = probabilities.getValues();
+
+    // The cdf won't include the last event. It will be implicit if not other
+    // event happened.
+    const cdf = new Float32Array(probabilities.size - 1);
+    cdf[0] = probVals[0];
+    for (let event = 1; event < cdf.length; ++event) {
+      cdf[event] = cdf[event - 1] + probVals[event];
+    }
+
+    const random = seedrandom(seed.toString());
+    const res = new Float32Array(numSamples);
+
+    for (let i = 0; i < numSamples; ++i) {
+      const r = random();
+
+      // Assume last event happened by default.
+      res[i] = cdf.length;
+
+      for (let event = 0; event < cdf.length; event++) {
+        if (r < cdf[event]) {
+          res[i] = event;
+          break;
+        }
+      }
+    }
+
+    return Array1D.new(res);
+  }
+
+  protected oneHotInternal(
+      indices: Array1D, depth: number, onValue: number,
+      offValue: number): Array2D {
+    const res = new Float32Array(indices.size * depth);
+    res.fill(offValue);
+
+    for (let event = 0; event < indices.size; ++event) {
+      res[event * depth + indices.get(event)] = onValue;
+    }
+    return Array2D.new([indices.size, depth], res);
+  }
 }

src/math/math_gpu.ts

@@ -38,6 +38,8 @@ import {LogSumExpProgram} from './webgl/logsumexp_gpu';
 import {MaxPool2DBackpropProgram} from './webgl/max_pool_backprop_gpu';
 import {MinMaxProgram} from './webgl/minmax_gpu';
 import {MatMulProgram} from './webgl/mulmat_gpu';
+import {MultinomialProgram} from './webgl/multinomial_gpu';
+import {OneHotProgram} from './webgl/onehot_gpu';
 import {Pool2DProgram} from './webgl/pool_gpu';
 import {ReduceSumProgram} from './webgl/reducesum_gpu';
 import {ResizeBilinear3DProgram} from './webgl/resize_bilinear_gpu';
@@ -424,6 +426,20 @@ export class NDArrayMathGPU extends NDArrayMath {
     return this.compileAndRun(program, [x]);
   }
 
+  protected multinomialInternal(
+      probs: Array1D, numSamples: number, seed: number): Array1D {
+    const program = new MultinomialProgram(probs.size, numSamples);
+    const customSetup = program.getCustomSetupFunc(seed);
+    return this.compileAndRun(program, [probs], null, customSetup);
+  }
+
+  protected oneHotInternal(
+      indices: ndarray.Array1D, depth: number, onValue: number,
+      offValue: number): ndarray.Array2D {
+    const program = new OneHotProgram(indices.size, depth, onValue, offValue);
+    return this.compileAndRun(program, [indices]);
+  }
+
   private getAndSaveBinary(key: string, getBinary: () => GPGPUBinary):
       GPGPUBinary {
     if (!(key in this.binaryCache)) {

src/math/multinomial_test.ts

@@ -0,0 +1,99 @@
+/**
+ * @license
+ * Copyright 2017 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 {NDArrayMath} from './math';
+import {NDArrayMathCPU} from './math_cpu';
+import {NDArrayMathGPU} from './math_gpu';
+import {Array1D} from './ndarray';
+
+function executeTests(mathFactory: () => NDArrayMath) {
+  let math: NDArrayMath;
+
+  beforeEach(() => {
+    math = mathFactory();
+    math.startScope();
+  });
+
+  afterEach(() => {
+    math.endScope(null);
+    math.dispose();
+  });
+
+  it('Flip a fair coin and check bounds', () => {
+    const probs = Array1D.new([0.5, 0.5]);
+    const result = math.multinomial(probs, 100);
+    expect(result.shape).toEqual([100]);
+    const [min, max] = getBounds(result.getValues());
+    expect(min >= 0 - 1e-4);
+    expect(max <= 1 + 1e-4);
+  });
+
+  it('Flip a two-sided coin with 100% of heads', () => {
+    const probs = Array1D.new([1, 0]);
+    const result = math.multinomial(probs, 100);
+    expect(result.shape).toEqual([100]);
+    const [min, max] = getBounds(result.getValues());
+    expect(min).toBeCloseTo(0, 1e-4);
+    expect(max).toBeCloseTo(0, 1e-4);
+  });
+
+  it('Flip a two-sided coin with 100% of tails', () => {
+    const probs = Array1D.new([0, 1]);
+    const result = math.multinomial(probs, 100);
+    expect(result.shape).toEqual([100]);
+    const [min, max] = getBounds(result.getValues());
+    expect(min).toBeCloseTo(1, 1e-4);
+    expect(max).toBeCloseTo(1, 1e-4);
+  });
+
+  it('Flip a single-sided coin throws error', () => {
+    const probs = Array1D.new([1]);
+    expect(() => math.multinomial(probs, 100)).toThrowError();
+  });
+
+  it('Flip a ten-sided coin and check bounds', () => {
+    const n = 10;
+    const probs = Array1D.zeros([n]);
+    for (let i = 0; i < n; ++i) {
+      probs.set(1 / n, i);
+    }
+    const result = math.multinomial(probs, 100);
+    expect(result.shape).toEqual([100]);
+    const [min, max] = getBounds(result.getValues());
+    expect(min >= 0 - 1e-4);
+    expect(max <= 9 + 1e-4);
+  });
+
+  function getBounds(a: Float32Array) {
+    let min = Number.MAX_VALUE;
+    let max = Number.MIN_VALUE;
+
+    for (let i = 0; i < a.length; ++i) {
+      min = Math.min(min, a[i]);
+      max = Math.max(max, a[i]);
+    }
+    return [min, max];
+  }
+}
+
+describe('mathCPU multinomial', () => {
+  executeTests(() => new NDArrayMathCPU());
+});
+
+describe('mathGPU multinomial', () => {
+  executeTests(() => new NDArrayMathGPU());
+});

src/math/onehot_test.ts

@@ -0,0 +1,81 @@
+/**
+ * @license
+ * Copyright 2017 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 test_util from '../test_util';
+import {NDArrayMath} from './math';
+import {NDArrayMathCPU} from './math_cpu';
+import {NDArrayMathGPU} from './math_gpu';
+import {Array1D} from './ndarray';
+
+function executeTests(mathFactory: () => NDArrayMath) {
+  let math: NDArrayMath;
+
+  beforeEach(() => {
+    math = mathFactory();
+    math.startScope();
+  });
+
+  afterEach(() => {
+    math.endScope(null);
+    math.dispose();
+  });
+
+  it('Depth 1 throws error', () => {
+    const indices = Array1D.new([0, 0, 0]);
+    expect(() => math.oneHot(indices, 1)).toThrowError();
+  });
+
+  it('Depth 2, diagonal', () => {
+    const indices = Array1D.new([0, 1]);
+    const res = math.oneHot(indices, 2);
+    const expected = new Float32Array([1, 0, 0, 1]);
+    expect(res.shape).toEqual([2, 2]);
+    test_util.expectArraysClose(res.getValues(), expected);
+  });
+
+  it('Depth 2, transposed diagonal', () => {
+    const indices = Array1D.new([1, 0]);
+    const res = math.oneHot(indices, 2);
+    const expected = new Float32Array([0, 1, 1, 0]);
+    expect(res.shape).toEqual([2, 2]);
+    test_util.expectArraysClose(res.getValues(), expected);
+  });
+
+  it('Depth 3, 4 events', () => {
+    const indices = Array1D.new([2, 1, 2, 0]);
+    const res = math.oneHot(indices, 3);
+    const expected = new Float32Array([0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0]);
+    expect(res.shape).toEqual([4, 3]);
+    test_util.expectArraysClose(res.getValues(), expected);
+  });
+
+  it('Depth 2 onValue=3, offValue=-2', () => {
+    const indices = Array1D.new([0, 1]);
+    const res = math.oneHot(indices, 2, 3, -2);
+    const expected = new Float32Array([3, -2, -2, 3]);
+    expect(res.shape).toEqual([2, 2]);
+    test_util.expectArraysClose(res.getValues(), expected);
+  });
+}
+
+describe('mathCPU oneHot', () => {
+  executeTests(() => new NDArrayMathCPU());
+});
+
+describe('mathGPU oneHot', () => {
+  executeTests(() => new NDArrayMathGPU());
+});