From 48780ef3a7d08dcb4945e846cc228f4869a030d9 Mon Sep 17 00:00:00 2001
From: Daniel Smilkov <dsmilkov@gmail.com>
Date: Thu, 28 Sep 2017 17:57:47 -0400
Subject: [PATCH 1/4] add multinomial

---
 src/math/math.ts                       |  15 ++++
 src/math/math_cpu.ts                   |   3 +-
 src/math/math_gpu.ts                   |   8 ++
 src/math/webgl/multinomial_gpu.ts      |  71 ++++++++++++++++
 src/math/webgl/multinomial_gpu_test.ts | 111 +++++++++++++++++++++++++
 5 files changed, 207 insertions(+), 1 deletion(-)
 create mode 100644 src/math/webgl/multinomial_gpu.ts
 create mode 100644 src/math/webgl/multinomial_gpu_test.ts

diff --git a/src/math/math.ts b/src/math/math.ts
index 904b63f5d1..7ce5e214da 100644
--- a/src/math/math.ts
+++ b/src/math/math.ts
@@ -1574,6 +1574,21 @@ 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 {
+    return this.track(
+        this.multinomialInternal(probabilities, numSamples, seed));
+  }
+  protected abstract multinomialInternal(
+      probabilities: Array1D, numSamples: number, seed: number): Array1D;
 }
 
 export enum MatrixOrientation {
diff --git a/src/math/math_cpu.ts b/src/math/math_cpu.ts
index 1558413b72..426d869f3c 100644
--- a/src/math/math_cpu.ts
+++ b/src/math/math_cpu.ts
@@ -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):
diff --git a/src/math/math_gpu.ts b/src/math/math_gpu.ts
index 5c2a2146f7..a13cb54fc2 100644
--- a/src/math/math_gpu.ts
+++ b/src/math/math_gpu.ts
@@ -38,6 +38,7 @@ 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 {Pool2DProgram} from './webgl/pool_gpu';
 import {ReduceSumProgram} from './webgl/reducesum_gpu';
 import {ResizeBilinear3DProgram} from './webgl/resize_bilinear_gpu';
@@ -424,6 +425,13 @@ 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 || Math.random());
+    return this.compileAndRun(program, [probs], null, customSetup);
+  }
+
   private getAndSaveBinary(key: string, getBinary: () => GPGPUBinary):
       GPGPUBinary {
     if (!(key in this.binaryCache)) {
diff --git a/src/math/webgl/multinomial_gpu.ts b/src/math/webgl/multinomial_gpu.ts
new file mode 100644
index 0000000000..b826f9e215
--- /dev/null
+++ b/src/math/webgl/multinomial_gpu.ts
@@ -0,0 +1,71 @@
+/**
+ * @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 {GPGPUContext} from './gpgpu_context';
+import {GPGPUProgram} from './gpgpu_math';
+
+export class MultinomialProgram implements GPGPUProgram {
+  variableNames = ['probs'];
+  params: Array<{}>;
+  outputShape: number[];
+  userCode: string;
+
+  // Caching uniform location for speed.
+  seedLoc: WebGLUniformLocation;
+
+  constructor(numOutcomes: number, numSamples: number) {
+    this.outputShape = [numSamples];
+    this.params = [];
+
+    this.userCode = `
+      uniform float seed;
+
+      const vec2 K1 = vec2(
+        23.14069263277926, // e^pi (Gelfond's constant)
+         2.665144142690225 // 2^sqrt(2) (Gelfond–Schneider constant)
+      );
+
+      float random(float seed) {
+          return fract(cos(dot(resultUV * seed, K1)) * 12345.6789);
+      }
+
+      void main() {
+        float r = random(seed);
+        float cdf = 0.0;
+
+        for (int i = 0; i < ${numOutcomes}; i++) {
+          cdf += getProbs(i);
+
+          if (r < cdf) {
+            setOutput(float(i));
+            return;
+          }
+        }
+        setOutput(float(${numOutcomes - 1}));
+      }
+    `;
+  }
+
+  getCustomSetupFunc(seed: number) {
+    return (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => {
+      if (this.seedLoc == null) {
+        this.seedLoc = gpgpu.getUniformLocation(webGLProgram, 'seed');
+      }
+      gpgpu.gl.uniform1f(this.seedLoc, seed);
+    };
+  }
+}
diff --git a/src/math/webgl/multinomial_gpu_test.ts b/src/math/webgl/multinomial_gpu_test.ts
new file mode 100644
index 0000000000..dc743d2f2a
--- /dev/null
+++ b/src/math/webgl/multinomial_gpu_test.ts
@@ -0,0 +1,111 @@
+/**
+ * @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 {Array1D, initializeGPU} from '../ndarray';
+import {GPGPUContext} from './gpgpu_context';
+import * as gpgpu_math from './gpgpu_math';
+import {MultinomialProgram} from './multinomial_gpu';
+import {TextureManager} from './texture_manager';
+
+describe('multinomial_gpu', () => {
+  let gpgpu: GPGPUContext;
+  let texManager: TextureManager;
+
+  beforeAll(() => {
+    gpgpu = new GPGPUContext();
+    texManager = new TextureManager(gpgpu);
+    initializeGPU(gpgpu, texManager);
+  });
+
+  afterAll(() => {
+    texManager.dispose();
+    gpgpu.dispose();
+  });
+
+  it('Flip a fair coin and check bounds', () => {
+    const probs = Array1D.new([0.5, 0.5]);
+    const result = doMultinomial(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 = doMultinomial(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 = doMultinomial(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', () => {
+    const probs = Array1D.new([1]);
+    const result = doMultinomial(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 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 = doMultinomial(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];
+  }
+
+  function doMultinomial(probs: Array1D, numSamples: number): Array1D {
+    const program = new MultinomialProgram(probs.size, numSamples);
+    const result = Array1D.zeros([numSamples]);
+
+    const binary = gpgpu_math.compileProgram(gpgpu, program, [probs], result);
+    const customSetup = program.getCustomSetupFunc(Math.random());
+    gpgpu_math.runProgram(binary, [probs], result, customSetup);
+
+    probs.dispose();
+    gpgpu.deleteProgram(binary.webGLProgram);
+    return result;
+  }
+});

From 5538d0471633da04d9a3f07c2d9acc837dc75e49 Mon Sep 17 00:00:00 2001
From: Daniel Smilkov <dsmilkov@gmail.com>
Date: Fri, 29 Sep 2017 11:32:58 -0400
Subject: [PATCH 2/4] add cpu implementation and tests

---
 package.json                                  |  4 ++
 src/math/math.ts                              | 17 ++++--
 src/math/math_cpu.ts                          | 34 ++++++++++-
 src/math/math_gpu.ts                          |  2 +-
 ...nomial_gpu_test.ts => multinomial_test.ts} | 59 ++++++++-----------
 src/math/webgl/multinomial_gpu.ts             |  4 +-
 6 files changed, 76 insertions(+), 44 deletions(-)
 rename src/math/{webgl/multinomial_gpu_test.ts => multinomial_test.ts} (59%)

diff --git a/package.json b/package.json
index fdd6e997ae..deea4df056 100644
--- a/package.json
+++ b/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"
   }
 }
diff --git a/src/math/math.ts b/src/math/math.ts
index 7ce5e214da..0be3d5a93d 100644
--- a/src/math/math.ts
+++ b/src/math/math.ts
@@ -270,9 +270,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();
   }
@@ -1584,8 +1583,16 @@ export abstract class NDArrayMath {
    */
   multinomial(probabilities: Array1D, numSamples: number, seed?: number):
       Array1D {
-    return this.track(
-        this.multinomialInternal(probabilities, numSamples, seed));
+    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;
diff --git a/src/math/math_cpu.ts b/src/math/math_cpu.ts
index 426d869f3c..b58ec03630 100644
--- a/src/math/math_cpu.ts
+++ b/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';
@@ -983,4 +983,36 @@ 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);
+  }
 }
diff --git a/src/math/math_gpu.ts b/src/math/math_gpu.ts
index a13cb54fc2..4855cd2b49 100644
--- a/src/math/math_gpu.ts
+++ b/src/math/math_gpu.ts
@@ -428,7 +428,7 @@ export class NDArrayMathGPU extends NDArrayMath {
   protected multinomialInternal(
       probs: Array1D, numSamples: number, seed: number): Array1D {
     const program = new MultinomialProgram(probs.size, numSamples);
-    const customSetup = program.getCustomSetupFunc(seed || Math.random());
+    const customSetup = program.getCustomSetupFunc(seed);
     return this.compileAndRun(program, [probs], null, customSetup);
   }
 
diff --git a/src/math/webgl/multinomial_gpu_test.ts b/src/math/multinomial_test.ts
similarity index 59%
rename from src/math/webgl/multinomial_gpu_test.ts
rename to src/math/multinomial_test.ts
index dc743d2f2a..80b4089bc7 100644
--- a/src/math/webgl/multinomial_gpu_test.ts
+++ b/src/math/multinomial_test.ts
@@ -15,30 +15,26 @@
  * =============================================================================
  */
 
-import {Array1D, initializeGPU} from '../ndarray';
-import {GPGPUContext} from './gpgpu_context';
-import * as gpgpu_math from './gpgpu_math';
-import {MultinomialProgram} from './multinomial_gpu';
-import {TextureManager} from './texture_manager';
+import {NDArrayMath} from './math';
+import {NDArrayMathCPU} from './math_cpu';
+import {NDArrayMathGPU} from './math_gpu';
+import {Array1D} from './ndarray';
 
-describe('multinomial_gpu', () => {
-  let gpgpu: GPGPUContext;
-  let texManager: TextureManager;
+function executeTests(mathFactory: () => NDArrayMath) {
+  let math: NDArrayMath;
 
-  beforeAll(() => {
-    gpgpu = new GPGPUContext();
-    texManager = new TextureManager(gpgpu);
-    initializeGPU(gpgpu, texManager);
+  beforeEach(() => {
+    math = mathFactory();
+    math.startScope();
   });
 
-  afterAll(() => {
-    texManager.dispose();
-    gpgpu.dispose();
+  afterEach(() => {
+    math.endScope(null);
   });
 
   it('Flip a fair coin and check bounds', () => {
     const probs = Array1D.new([0.5, 0.5]);
-    const result = doMultinomial(probs, 100);
+    const result = math.multinomial(probs, 100);
     expect(result.shape).toEqual([100]);
     const [min, max] = getBounds(result.getValues());
     expect(min >= 0 - 1e-4);
@@ -47,7 +43,7 @@ describe('multinomial_gpu', () => {
 
   it('Flip a two-sided coin with 100% of heads', () => {
     const probs = Array1D.new([1, 0]);
-    const result = doMultinomial(probs, 100);
+    const result = math.multinomial(probs, 100);
     expect(result.shape).toEqual([100]);
     const [min, max] = getBounds(result.getValues());
     expect(min).toBeCloseTo(0, 1e-4);
@@ -56,20 +52,16 @@ describe('multinomial_gpu', () => {
 
   it('Flip a two-sided coin with 100% of tails', () => {
     const probs = Array1D.new([0, 1]);
-    const result = doMultinomial(probs, 100);
+    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', () => {
+  it('Flip a single-sided coin throws error', () => {
     const probs = Array1D.new([1]);
-    const result = doMultinomial(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);
+    expect(() => math.multinomial(probs, 100)).toThrowError();
   });
 
   it('Flip a ten-sided coin and check bounds', () => {
@@ -78,7 +70,7 @@ describe('multinomial_gpu', () => {
     for (let i = 0; i < n; ++i) {
       probs.set(1 / n, i);
     }
-    const result = doMultinomial(probs, 100);
+    const result = math.multinomial(probs, 100);
     expect(result.shape).toEqual([100]);
     const [min, max] = getBounds(result.getValues());
     expect(min >= 0 - 1e-4);
@@ -95,17 +87,12 @@ describe('multinomial_gpu', () => {
     }
     return [min, max];
   }
+}
 
-  function doMultinomial(probs: Array1D, numSamples: number): Array1D {
-    const program = new MultinomialProgram(probs.size, numSamples);
-    const result = Array1D.zeros([numSamples]);
-
-    const binary = gpgpu_math.compileProgram(gpgpu, program, [probs], result);
-    const customSetup = program.getCustomSetupFunc(Math.random());
-    gpgpu_math.runProgram(binary, [probs], result, customSetup);
+describe('mathCPU.multinomial', () => {
+  executeTests(() => new NDArrayMathCPU());
+});
 
-    probs.dispose();
-    gpgpu.deleteProgram(binary.webGLProgram);
-    return result;
-  }
+describe('mathGPU.multinomial', () => {
+  executeTests(() => new NDArrayMathGPU());
 });
diff --git a/src/math/webgl/multinomial_gpu.ts b/src/math/webgl/multinomial_gpu.ts
index b826f9e215..7c895243c1 100644
--- a/src/math/webgl/multinomial_gpu.ts
+++ b/src/math/webgl/multinomial_gpu.ts
@@ -47,7 +47,7 @@ export class MultinomialProgram implements GPGPUProgram {
         float r = random(seed);
         float cdf = 0.0;
 
-        for (int i = 0; i < ${numOutcomes}; i++) {
+        for (int i = 0; i < ${numOutcomes - 1}; i++) {
           cdf += getProbs(i);
 
           if (r < cdf) {
@@ -55,6 +55,8 @@ export class MultinomialProgram implements GPGPUProgram {
             return;
           }
         }
+
+        // If no other event happened, last event happened.
         setOutput(float(${numOutcomes - 1}));
       }
     `;

From 1fe0684e1595dfed6fd589217c19f386a2033974 Mon Sep 17 00:00:00 2001
From: Daniel Smilkov <dsmilkov@gmail.com>
Date: Fri, 29 Sep 2017 17:31:20 -0400
Subject: [PATCH 3/4] add onehot operation

---
 src/math/math.ts                  | 30 +++++++++++++
 src/math/math_cpu.ts              | 11 +++++
 src/math/math_gpu.ts              |  8 ++++
 src/math/multinomial_test.ts      |  4 +-
 src/math/onehot_test.ts           | 72 +++++++++++++++++++++++++++++++
 src/math/webgl/onehot_gpu.ts      | 53 +++++++++++++++++++++++
 src/math/webgl/shader_compiler.ts |  4 ++
 7 files changed, 180 insertions(+), 2 deletions(-)
 create mode 100644 src/math/onehot_test.ts
 create mode 100644 src/math/webgl/onehot_gpu.ts

diff --git a/src/math/math.ts b/src/math/math.ts
index 0be3d5a93d..7e7292f374 100644
--- a/src/math/math.ts
+++ b/src/math/math.ts
@@ -1596,6 +1596,36 @@ export abstract class NDArrayMath {
   }
   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?: number, offValue?: number):
+      Array2D {
+    if (depth < 2) {
+      throw new Error(`Error in oneHot: depth must be >=2, but it is ${depth}`);
+    }
+    if (onValue == null) {
+      onValue = 1;
+    }
+    if (offValue == null) {
+      offValue = 0;
+    }
+    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 {
diff --git a/src/math/math_cpu.ts b/src/math/math_cpu.ts
index b58ec03630..c9d87eb93a 100644
--- a/src/math/math_cpu.ts
+++ b/src/math/math_cpu.ts
@@ -1015,4 +1015,15 @@ export class NDArrayMathCPU extends NDArrayMath {
 
     return Array1D.new(res);
   }
+
+  protected oneHotInternal(
+      indices: Array1D, depth: number, onValue: number,
+      offValue: number): Array2D {
+    const res = new Float32Array(indices.size * depth);
+
+    for (let event = 0; event < indices.size; ++event) {
+      res[event * depth + indices.get(event)] = 1;
+    }
+    return Array2D.new([indices.size, depth], res);
+  }
 }
diff --git a/src/math/math_gpu.ts b/src/math/math_gpu.ts
index 4855cd2b49..c55a8bccee 100644
--- a/src/math/math_gpu.ts
+++ b/src/math/math_gpu.ts
@@ -39,6 +39,7 @@ 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';
@@ -432,6 +433,13 @@ export class NDArrayMathGPU extends NDArrayMath {
     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)) {
diff --git a/src/math/multinomial_test.ts b/src/math/multinomial_test.ts
index 80b4089bc7..973a162599 100644
--- a/src/math/multinomial_test.ts
+++ b/src/math/multinomial_test.ts
@@ -89,10 +89,10 @@ function executeTests(mathFactory: () => NDArrayMath) {
   }
 }
 
-describe('mathCPU.multinomial', () => {
+describe('mathCPU multinomial', () => {
   executeTests(() => new NDArrayMathCPU());
 });
 
-describe('mathGPU.multinomial', () => {
+describe('mathGPU multinomial', () => {
   executeTests(() => new NDArrayMathGPU());
 });
diff --git a/src/math/onehot_test.ts b/src/math/onehot_test.ts
new file mode 100644
index 0000000000..18155f4112
--- /dev/null
+++ b/src/math/onehot_test.ts
@@ -0,0 +1,72 @@
+/**
+ * @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);
+  });
+
+  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);
+  });
+}
+
+describe('mathCPU oneHot', () => {
+  executeTests(() => new NDArrayMathCPU());
+});
+
+describe('mathGPU oneHot', () => {
+  executeTests(() => new NDArrayMathGPU());
+});
diff --git a/src/math/webgl/onehot_gpu.ts b/src/math/webgl/onehot_gpu.ts
new file mode 100644
index 0000000000..43c78d2003
--- /dev/null
+++ b/src/math/webgl/onehot_gpu.ts
@@ -0,0 +1,53 @@
+/**
+ * @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 {GPGPUContext} from './gpgpu_context';
+import {GPGPUProgram} from './gpgpu_math';
+
+export class OneHotProgram implements GPGPUProgram {
+  variableNames = ['indices'];
+  params: Array<{}>;
+  outputShape: number[];
+  userCode: string;
+
+  // Caching uniform location for speed.
+  seedLoc: WebGLUniformLocation;
+
+  constructor(
+      numIndices: number, depth: number, onValue: number, offValue: number) {
+    this.outputShape = [numIndices, depth];
+    this.params = [onValue, offValue];
+
+    this.userCode = `
+      void main() {
+        ivec2 coords = getOutputCoords();
+        int index = round(getIndices(coords.x));
+        setOutput(mix(float(${offValue}), float(${onValue}),
+                      float(index == coords.y)));
+      }
+    `;
+  }
+
+  getCustomSetupFunc(seed: number) {
+    return (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => {
+      if (this.seedLoc == null) {
+        this.seedLoc = gpgpu.getUniformLocation(webGLProgram, 'seed');
+      }
+      gpgpu.gl.uniform1f(this.seedLoc, seed);
+    };
+  }
+}
diff --git a/src/math/webgl/shader_compiler.ts b/src/math/webgl/shader_compiler.ts
index 5da72bc0ea..dbc613ad91 100644
--- a/src/math/webgl/shader_compiler.ts
+++ b/src/math/webgl/shader_compiler.ts
@@ -173,6 +173,10 @@ const SHADER_PREFIX = `
     return dot(vec4(1), values);
   }
 
+  int round(float value) {
+    return int(floor(value + 0.5));
+  }
+
   ${SAMPLE_1D_SNIPPET}
   ${SAMPLE_2D_SNIPPET}
   ${SAMPLE_3D_SNIPPET}

From 32f9b6ee57d80fe89e41e6daae48a42ffdddbcc7 Mon Sep 17 00:00:00 2001
From: Daniel Smilkov <dsmilkov@gmail.com>
Date: Sat, 30 Sep 2017 10:21:20 -0400
Subject: [PATCH 4/4] resolve comments

---
 package.json                      |  2 +-
 src/math/math.ts                  | 12 ++++--------
 src/math/math_cpu.ts              |  3 ++-
 src/math/multinomial_test.ts      |  1 +
 src/math/onehot_test.ts           |  9 +++++++++
 src/math/webgl/multinomial_gpu.ts |  9 ---------
 src/math/webgl/shader_compiler.ts |  9 +++++++++
 7 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/package.json b/package.json
index deea4df056..c8886a5659 100644
--- a/package.json
+++ b/package.json
@@ -13,7 +13,7 @@
   "devDependencies": {
     "@types/jasmine": "~2.5.53",
     "@types/polymer": "~1.1.31",
-    "@types/seedrandom": "^2.4.27",
+    "@types/seedrandom": "~2.4.27",
     "bower": "~1.8.0",
     "browserify": "~14.4.0",
     "cross-spawn": "~5.1.0",
diff --git a/src/math/math.ts b/src/math/math.ts
index 7e7292f374..82f0527164 100644
--- a/src/math/math.ts
+++ b/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
@@ -1609,17 +1612,10 @@ export abstract class NDArrayMath {
    * @param offValue A number used to fill in the output when the index does not
    *     match the location.
    */
-  oneHot(indices: Array1D, depth: number, onValue?: number, offValue?: number):
-      Array2D {
+  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}`);
     }
-    if (onValue == null) {
-      onValue = 1;
-    }
-    if (offValue == null) {
-      offValue = 0;
-    }
     return this.executeOp(
         'oneHot', () => this.oneHotInternal(indices, depth, onValue, offValue));
   }
diff --git a/src/math/math_cpu.ts b/src/math/math_cpu.ts
index c9d87eb93a..3b6f3d73fb 100644
--- a/src/math/math_cpu.ts
+++ b/src/math/math_cpu.ts
@@ -1020,9 +1020,10 @@ export class NDArrayMathCPU extends NDArrayMath {
       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)] = 1;
+      res[event * depth + indices.get(event)] = onValue;
     }
     return Array2D.new([indices.size, depth], res);
   }
diff --git a/src/math/multinomial_test.ts b/src/math/multinomial_test.ts
index 973a162599..81a0aa036b 100644
--- a/src/math/multinomial_test.ts
+++ b/src/math/multinomial_test.ts
@@ -30,6 +30,7 @@ function executeTests(mathFactory: () => NDArrayMath) {
 
   afterEach(() => {
     math.endScope(null);
+    math.dispose();
   });
 
   it('Flip a fair coin and check bounds', () => {
diff --git a/src/math/onehot_test.ts b/src/math/onehot_test.ts
index 18155f4112..83c850d443 100644
--- a/src/math/onehot_test.ts
+++ b/src/math/onehot_test.ts
@@ -31,6 +31,7 @@ function executeTests(mathFactory: () => NDArrayMath) {
 
   afterEach(() => {
     math.endScope(null);
+    math.dispose();
   });
 
   it('Depth 1 throws error', () => {
@@ -61,6 +62,14 @@ function executeTests(mathFactory: () => NDArrayMath) {
     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', () => {
diff --git a/src/math/webgl/multinomial_gpu.ts b/src/math/webgl/multinomial_gpu.ts
index 7c895243c1..19a58e8548 100644
--- a/src/math/webgl/multinomial_gpu.ts
+++ b/src/math/webgl/multinomial_gpu.ts
@@ -34,15 +34,6 @@ export class MultinomialProgram implements GPGPUProgram {
     this.userCode = `
       uniform float seed;
 
-      const vec2 K1 = vec2(
-        23.14069263277926, // e^pi (Gelfond's constant)
-         2.665144142690225 // 2^sqrt(2) (Gelfond–Schneider constant)
-      );
-
-      float random(float seed) {
-          return fract(cos(dot(resultUV * seed, K1)) * 12345.6789);
-      }
-
       void main() {
         float r = random(seed);
         float cdf = 0.0;
diff --git a/src/math/webgl/shader_compiler.ts b/src/math/webgl/shader_compiler.ts
index dbc613ad91..b8c9cdff97 100644
--- a/src/math/webgl/shader_compiler.ts
+++ b/src/math/webgl/shader_compiler.ts
@@ -177,6 +177,15 @@ const SHADER_PREFIX = `
     return int(floor(value + 0.5));
   }
 
+  const vec2 randomConst = vec2(
+    23.14069263277926, // e^pi (Gelfond's constant)
+     2.665144142690225 // 2^sqrt(2) (Gelfond–Schneider constant)
+  );
+
+  float random(float seed) {
+      return fract(cos(dot(resultUV * seed, randomConst)) * 12345.6789);
+  }
+
   ${SAMPLE_1D_SNIPPET}
   ${SAMPLE_2D_SNIPPET}
   ${SAMPLE_3D_SNIPPET}