From 1555ed84723cc4cd38dd6a6a0e13ebd28160a53d Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 09:35:23 -0500
Subject: [PATCH 01/84] clean

---
 tfjs-core/src/backends/backend.ts             |  3 ++
 tfjs-core/src/backends/webgl/backend_webgl.ts |  5 +++
 tfjs-core/src/ops/softmax.ts                  | 41 +++++++++++--------
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts
index fd3dbbaab3d..9901a3614db 100644
--- a/tfjs-core/src/backends/backend.ts
+++ b/tfjs-core/src/backends/backend.ts
@@ -305,6 +305,9 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
   expm1<T extends Tensor>(x: T): T {
     return notYetImplemented('expm1');
   }
+  softmax<T extends Tensor>(x: T): T {
+    return notYetImplemented('softmax');
+  }
   log<T extends Tensor>(x: T): T {
     return notYetImplemented('log');
   }
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 2bce838b3aa..d62353d4a4b 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1588,6 +1588,11 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
+  softmax<T extends Tensor>(x: T): T {
+    // const program = new
+    return x;
+  }
+
   log<T extends Tensor>(x: T): T {
     if (this.shouldExecuteOnCPU([x])) {
       return this.cpuBackend.log(x);
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index c308db2ded4..acc72761773 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -15,13 +15,16 @@
  * =============================================================================
  */
 
+import {ENGINE} from '../engine';
 import {customGrad} from '../gradients';
 import {Tensor} from '../tensor';
 import {GradSaveFunc} from '../tensor_types';
 import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
+
 import {op} from './operation';
 
+
 /**
  * Computes the softmax normalized vector given the logits.
  *
@@ -54,25 +57,29 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
         `Logits was rank ${$logits.rank} and dim was ${dim}`);
   }
 
-  const customOp = customGrad((logits: Tensor, save: GradSaveFunc) => {
-    // Do it in log space for numerical stability.
-    // exp(X - logSumExp(X))
-    const keepDims = true;
-    const lse = logits.logSumExp([dim], keepDims);
-    const logResult = logits.toFloat().sub(lse);
-    const y = logResult.exp() as T;
-    save([y]);
-    const gradFunc = (dy: T, saved: Tensor[]) => {
-      const [y] = saved;
-      const dyTimesY = dy.mul(y);
-      const keepDims = true;
-      return dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y));
-    };
+  return ENGINE.runKernelFunc(
+      (_, save) => {
+        const keepDims = true;
+        const lse = $logits.logSumExp([dim], keepDims);
+        const logResult = $logits.toFloat().sub(lse);
+        const y = logResult.exp() as T;
+        save([y]);
 
-    return {value: y, gradFunc};
-  });
+        return y;
+      },
+      {logits: $logits},
+      (dy: T, saved: Tensor[]) => {
+        const [y] = saved;
+        const dyTimesY = dy.mul(y);
+        const keepDims = true;
 
-  return customOp($logits);
+        return {
+          logits: () => {
+            return dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y));
+          }
+        };
+      },
+      'Softmax', {dim});
 }
 
 /**

From 22171d3fccd1141912939270f5456bd69bdf86e4 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 09:50:08 -0500
Subject: [PATCH 02/84] shell

---
 tfjs-core/src/backends/webgl/backend_webgl.ts |  4 ++--
 tfjs-core/src/ops/softmax.ts                  | 15 ++++++++++-----
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index d62353d4a4b..f4af2474083 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1589,8 +1589,8 @@ export class MathBackendWebGL extends KernelBackend {
   }
 
   softmax<T extends Tensor>(x: T): T {
-    // const program = new
-    return x;
+    const program = new UnaryOpProgram(x.shape, unary_op.EXP);
+    return this.compileAndRun(program, [x]);
   }
 
   log<T extends Tensor>(x: T): T {
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index acc72761773..a4b9dd650bd 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -58,11 +58,16 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
   }
 
   return ENGINE.runKernelFunc(
-      (_, save) => {
-        const keepDims = true;
-        const lse = $logits.logSumExp([dim], keepDims);
-        const logResult = $logits.toFloat().sub(lse);
-        const y = logResult.exp() as T;
+      (backend, save) => {
+        // SAVED WORKING IMPLEMENTATION
+        // const keepDims = true;
+        // const lse = $logits.logSumExp([dim], keepDims);
+        // const logResult = $logits.toFloat().sub(lse);
+        // const y = logResult.exp() as T;
+
+        // NEW KERNEL
+        const y = backend.softmax($logits);
+
         save([y]);
 
         return y;

From 99b1e5c7577d43a6b4a08a1f35769b19fd6b451b Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 10:34:20 -0500
Subject: [PATCH 03/84] compute lse

---
 tfjs-core/src/backends/webgl/backend_webgl.ts | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index f4af2474083..721739a123f 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1588,9 +1588,21 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  softmax<T extends Tensor>(x: T): T {
-    const program = new UnaryOpProgram(x.shape, unary_op.EXP);
-    return this.compileAndRun(program, [x]);
+  softmax<T extends Tensor>(logits: T): T {
+    const axes = util.parseAxisParam([logits.rank - 1], logits.shape);
+    const maxLogit = this.max(logits, axes);
+    const a = this.subtract(
+        logits,
+        maxLogit.reshape(axis_util.expandShapeToKeepDim(maxLogit.shape, axes)));
+    const b = this.exp(a);
+    const c = this.sum(b, axes);
+    const d = this.log(c);
+    let lse = this.add(this.reshape(maxLogit, d.shape), d);
+    const newShape = axis_util.expandShapeToKeepDim(lse.shape, axes);
+    lse = this.reshape(lse, newShape);
+
+    const program = new UnaryOpProgram(logits.shape, unary_op.EXP);
+    return this.compileAndRun(program, [logits]);
   }
 
   log<T extends Tensor>(x: T): T {

From 3865aac1957634a26e0bfa91fafff28177bb3856 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 10:41:36 -0500
Subject: [PATCH 04/84] softmax

---
 tfjs-core/src/backends/webgl/backend_webgl.ts |  5 +++--
 tfjs-core/src/ops/softmax.ts                  | 11 +----------
 2 files changed, 4 insertions(+), 12 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 721739a123f..1ca7e956e0a 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1601,8 +1601,9 @@ export class MathBackendWebGL extends KernelBackend {
     const newShape = axis_util.expandShapeToKeepDim(lse.shape, axes);
     lse = this.reshape(lse, newShape);
 
-    const program = new UnaryOpProgram(logits.shape, unary_op.EXP);
-    return this.compileAndRun(program, [logits]);
+    const logResult = this.subtract(logits, lse);
+    const y = this.exp(logResult);
+    return y as T;
   }
 
   log<T extends Tensor>(x: T): T {
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index a4b9dd650bd..f5f6e5e20ab 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -46,7 +46,7 @@ import {op} from './operation';
  */
 /** @doc {heading: 'Operations', subheading: 'Normalization'} */
 function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
-  const $logits = convertToTensor(logits, 'logits', 'softmax');
+  const $logits = convertToTensor(logits, 'logits', 'softmax', 'float32');
 
   if (dim === -1) {
     dim = $logits.rank - 1;
@@ -59,17 +59,8 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
 
   return ENGINE.runKernelFunc(
       (backend, save) => {
-        // SAVED WORKING IMPLEMENTATION
-        // const keepDims = true;
-        // const lse = $logits.logSumExp([dim], keepDims);
-        // const logResult = $logits.toFloat().sub(lse);
-        // const y = logResult.exp() as T;
-
-        // NEW KERNEL
         const y = backend.softmax($logits);
-
         save([y]);
-
         return y;
       },
       {logits: $logits},

From 971fe5a91f96f0d29d4be811843b69fd0f3da5de Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 10:48:36 -0500
Subject: [PATCH 05/84] add cpu

---
 tfjs-core/src/backends/cpu/backend_cpu.ts | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index ee2a1e6772a..7f9dbd0e620 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,6 +377,24 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
+  softmax<T extends Tensor>(logits: T): T {
+    const axes = util.parseAxisParam([logits.rank - 1], logits.shape);
+    const maxLogit = this.max(logits, axes);
+    const a = this.subtract(
+        logits,
+        maxLogit.reshape(axis_util.expandShapeToKeepDim(maxLogit.shape, axes)));
+    const b = this.exp(a);
+    const c = this.sum(b, axes);
+    const d = this.log(c);
+    let lse = this.add(this.reshape(maxLogit, d.shape), d);
+    const newShape = axis_util.expandShapeToKeepDim(lse.shape, axes);
+    lse = this.reshape(lse, newShape);
+
+    const logResult = this.subtract(logits, lse);
+    const y = this.exp(logResult);
+    return y as T;
+  }
+
   subtract(a: Tensor, b: Tensor): Tensor {
     if (a.dtype === 'complex64' || b.dtype === 'complex64') {
       return this.broadcastedBinaryComplexOp(

From 14100b7f7d8cacbf4a826bd8da9eb52c0a1278b8 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 11:30:27 -0500
Subject: [PATCH 06/84] upgrade deps

---
 tfjs-backend-wasm/WORKSPACE    | 30 ++++++++++++++++--------------
 tfjs-backend-wasm/package.json |  4 ++--
 tfjs-backend-wasm/yarn.lock    | 14 +++-----------
 tfjs-core/src/ops/softmax.ts   |  4 +---
 4 files changed, 22 insertions(+), 30 deletions(-)

diff --git a/tfjs-backend-wasm/WORKSPACE b/tfjs-backend-wasm/WORKSPACE
index 540effdb753..88f02bc54f4 100644
--- a/tfjs-backend-wasm/WORKSPACE
+++ b/tfjs-backend-wasm/WORKSPACE
@@ -8,78 +8,80 @@ emsdk_configure(name = "emsdk")
 
 git_repository(
     name = "xnnpack",
-    commit = "3a77ea7bbe30b2411591f2ab15f9c5032a25f688",
+    commit = "7278a95e3cfae6eac73f363c4fda5db53e1b2a87",
     remote = "https://github.com/google/XNNPACK.git",
-    shallow_since = "1577131863 -0800",
+    shallow_since = "1580796377 -0800",
 )
 
 # The libraries below are transitive dependencies of XNNPACK that we need to
 # explicitly enumerate here. See https://docs.bazel.build/versions/master/external.html#transitive-dependencies
 
+
 # FP16 library, used for half-precision conversions
 http_archive(
     name = "FP16",
-    build_file = "@xnnpack//third_party:FP16.BUILD",
-    sha256 = "9764297a339ad73b0717331a2c3e9c42a52105cd04cab62cb160e2b4598d2ea6",
     strip_prefix = "FP16-ba1d31f5eed2eb4a69e4dea3870a68c7c95f998f",
+    sha256 = "9764297a339ad73b0717331a2c3e9c42a52105cd04cab62cb160e2b4598d2ea6",
     urls = [
         "https://github.com/Maratyszcza/FP16/archive/ba1d31f5eed2eb4a69e4dea3870a68c7c95f998f.tar.gz",
     ],
+    build_file = "@xnnpack//third_party:FP16.BUILD",
 )
 
 # FXdiv library, used for repeated integer division by the same factor
 http_archive(
     name = "FXdiv",
-    build_file = "@xnnpack//third_party:FXdiv.BUILD",
-    sha256 = "7d3215bea832fe77091ec5666200b91156df6724da1e348205078346325fc45e",
     strip_prefix = "FXdiv-f8c5354679ec2597792bc70a9e06eff50c508b9a",
+    sha256 = "7d3215bea832fe77091ec5666200b91156df6724da1e348205078346325fc45e",
     urls = [
         "https://github.com/Maratyszcza/FXdiv/archive/f8c5354679ec2597792bc70a9e06eff50c508b9a.tar.gz",
     ],
+    build_file = "@xnnpack//third_party:FXdiv.BUILD",
 )
 
 # pthreadpool library, used for parallelization
 http_archive(
     name = "pthreadpool",
-    build_file = "@xnnpack//third_party:pthreadpool.BUILD",
     sha256 = "c2328fdf9e48ac9b928953bcbc442eb14402d393e4cfae0541581a3d39efca9d",
     strip_prefix = "pthreadpool-0e275fe56094626349c55a524ea8b71a85daa64b",
     urls = [
-        "https://github.com/Maratyszcza/pthreadpool/archive/0e275fe56094626349c55a524ea8b71a85daa64b.tar.gz",
+      "https://github.com/Maratyszcza/pthreadpool/archive/0e275fe56094626349c55a524ea8b71a85daa64b.tar.gz",
     ],
+    build_file = "@xnnpack//third_party:pthreadpool.BUILD",
 )
 
 # clog library, used for logging
 http_archive(
     name = "clog",
-    build_file = "@xnnpack//third_party:clog.BUILD",
-    sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
     strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
+    sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
     urls = [
         "https://github.com/pytorch/cpuinfo/archive/d5e37adf1406cf899d7d9ec1d317c47506ccb970.tar.gz",
     ],
+    build_file = "@xnnpack//third_party:clog.BUILD",
 )
 
 # cpuinfo library, used for detecting processor characteristics
 http_archive(
     name = "cpuinfo",
-    build_file = "@xnnpack//third_party:cpuinfo.BUILD",
-    sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
     strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
+    sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
     urls = [
         "https://github.com/pytorch/cpuinfo/archive/d5e37adf1406cf899d7d9ec1d317c47506ccb970.tar.gz",
     ],
+    build_file = "@xnnpack//third_party:cpuinfo.BUILD",
+    patches = ["@xnnpack//third_party:cpuinfo.patch"],
 )
 
 # psimd library, used for fallback 128-bit SIMD micro-kernels
 http_archive(
     name = "psimd",
-    build_file = "@xnnpack//third_party:psimd.BUILD",
-    sha256 = "1fefd66702cb2eb3462b962f33d4fb23d59a55d5889ee6372469d286c4512df4",
     strip_prefix = "psimd-10b4ffc6ea9e2e11668f86969586f88bc82aaefa",
+    sha256 = "1fefd66702cb2eb3462b962f33d4fb23d59a55d5889ee6372469d286c4512df4",
     urls = [
         "https://github.com/Maratyszcza/psimd/archive/10b4ffc6ea9e2e11668f86969586f88bc82aaefa.tar.gz",
     ],
+    build_file = "@xnnpack//third_party:psimd.BUILD",
 )
 
 git_repository(
diff --git a/tfjs-backend-wasm/package.json b/tfjs-backend-wasm/package.json
index e2fc73893a5..d86c9f17bc1 100644
--- a/tfjs-backend-wasm/package.json
+++ b/tfjs-backend-wasm/package.json
@@ -32,7 +32,7 @@
     "path": false
   },
   "peerDependencies": {
-    "@tensorflow/tfjs-core": "1.5.2"
+    "@tensorflow/tfjs-core": "link:../tfjs-core"
   },
   "dependencies": {
     "@types/emscripten": "~0.0.34"
@@ -40,7 +40,7 @@
   "devDependencies": {
     "@bazel/bazel": "^0.28.0",
     "@bazel/buildifier": "0.29.0",
-    "@tensorflow/tfjs-core": "1.5.2",
+    "@tensorflow/tfjs-core": "link:../tfjs-core",
     "@types/jasmine": "~2.8.6",
     "clang-format": "~1.2.4",
     "jasmine": "~3.1.0",
diff --git a/tfjs-backend-wasm/yarn.lock b/tfjs-backend-wasm/yarn.lock
index 2e76213ffed..ac1064763a2 100644
--- a/tfjs-backend-wasm/yarn.lock
+++ b/tfjs-backend-wasm/yarn.lock
@@ -73,17 +73,9 @@
   resolved "https://registry.yarnpkg.com/@bazel/hide-bazel-files/-/hide-bazel-files-0.38.3.tgz#e98231d3d360d51860d9c1a7c3345b40dab4cf81"
   integrity sha512-o+dNkfDm3qxWQ8h/04cWuTcjR7qnjZi3pQGv4aklVb16oPWx2jF8BzbkwvWuIkdbOl9VnqYP0vaHzwQVJRRcIA==
 
-"@tensorflow/tfjs-core@1.5.2":
-  version "1.5.2"
-  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-core/-/tfjs-core-1.5.2.tgz#df76752cf7c43987df1548fb69820935bd8215d7"
-  integrity sha512-Rj6l8xf0PxrEKctvX3bvxjqhHLaCBQT0ChvqFK6//HBu8A1/ao4SzeVKpXKNnP9Niax+qV3c9U9VcOwwIkCMag==
-  dependencies:
-    "@types/offscreencanvas" "~2019.3.0"
-    "@types/seedrandom" "2.4.27"
-    "@types/webgl-ext" "0.0.30"
-    "@types/webgl2" "0.0.4"
-    node-fetch "~2.1.2"
-    seedrandom "2.4.3"
+"@tensorflow/tfjs-core@link:../tfjs-core":
+  version "0.0.0"
+  uid ""
 
 "@types/emscripten@~0.0.34":
   version "0.0.34"
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index f5f6e5e20ab..b9144cf2af0 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -70,9 +70,7 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
         const keepDims = true;
 
         return {
-          logits: () => {
-            return dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y));
-          }
+          logits: () => dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y))
         };
       },
       'Softmax', {dim});

From cb717afebadba5d6356b9b51de0ee61a3c93df22 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 13:50:39 -0500
Subject: [PATCH 07/84] setup

---
 tfjs-backend-wasm/src/cc/BUILD                | 18 ++++
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc   | 99 +++++++++++++++++++
 .../src/cc/kernels/Softmax_test.cc            | 59 +++++++++++
 tfjs-backend-wasm/src/index_test.ts           | 15 ++-
 tfjs-backend-wasm/src/kernels/Softmax.ts      | 54 ++++++++++
 tfjs-backend-wasm/src/kernels/all_kernels.ts  |  1 +
 tfjs-backend-wasm/src/setup_test.ts           |  1 +
 tfjs-core/src/ops/softmax.ts                  |  2 +-
 8 files changed, 246 insertions(+), 3 deletions(-)
 create mode 100644 tfjs-backend-wasm/src/cc/kernels/Softmax.cc
 create mode 100644 tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
 create mode 100644 tfjs-backend-wasm/src/kernels/Softmax.ts

diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index d41cca267f1..4209ce300b6 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -186,6 +186,7 @@ tfjs_cc_library(
         ":ResizeBilinear",
         ":ScatterNd",
         ":Sigmoid",
+        ":Softmax",
         ":Sub",
         ":Tile",
         ":Transpose",
@@ -647,6 +648,15 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "Softmax",
+    srcs = ["kernels/Softmax.cc"],
+    deps = [
+        ":backend",
+        ":unary",
+    ],
+)
+
 tfjs_unit_test(
     name = "Sigmoid_test",
     srcs = ["kernels/Sigmoid_test.cc"],
@@ -655,6 +665,14 @@ tfjs_unit_test(
     ],
 )
 
+tfjs_unit_test(
+    name = "Softmax_test",
+    srcs = ["kernels/Softmax_test.cc"],
+    deps = [
+        ":Softmax",
+    ],
+)
+
 tfjs_cc_library(
     name = "Square",
     srcs = ["kernels/Square.cc"],
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
new file mode 100644
index 00000000000..75c0b234ca8
--- /dev/null
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -0,0 +1,99 @@
+/* Copyright 2020 Google LLC. 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <xnnpack.h>
+#include <cmath>
+#include <cstddef>
+#include <map>
+#include <tuple>
+#include <vector>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+// We use std::tuple as the cache key as it implements the compare operator
+// needed for std::map.
+typedef std::tuple<size_t> OperatorCacheKey;
+
+// The operator cache maps the weights id to the xnn_operator_t instantiated for
+// this set of weights.
+std::map<OperatorCacheKey, xnn_operator_t> operator_cache;
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+
+void Softmax(size_t x_id, size_t out_id) {
+  auto& x_info = backend::get_tensor_info(x_id);
+  auto& out_info = backend::get_tensor_info_out(out_id);
+
+  const float* x_buf = x_info.f32();
+  float* out_buf = out_info.f32_write();
+
+  xnn_operator_t softmax_op = nullptr;
+
+  const size_t channels = 3;
+  OperatorCacheKey cache_key = {channels};
+
+  auto operator_cache_idx = operator_cache.find(cache_key);
+  if (operator_cache_idx == operator_cache.end()) {
+    const size_t input_stride = 1;
+    const size_t output_stride = 1;
+    const uint32_t flags = 0;
+
+    xnn_status status = xnn_create_softmax_nc_f32(
+        channels, input_stride, output_stride, flags, &softmax_op);
+    if (status != xnn_status_success) {
+      tfjs::util::warn(
+          "XNN status for xnn_create_softmax_nc_f32 is not "
+          "successful. Got status %d. Use -c dbg to see XNN logs.",
+          status);
+      return;
+    }
+
+    operator_cache.insert({cache_key, softmax_op});
+
+    tfjs::backend::xnn_operator_count++;
+  } else {
+    softmax_op = operator_cache_idx->second;
+  }
+
+  xnn_status status = xnn_setup_softmax_nc_f32(softmax_op, 1, x_buf, out_buf,
+                                               nullptr /* thread pool */);
+  if (status != xnn_status_success) {
+    tfjs::util::warn(
+        "XNN status for xnn_setup_softmax_nc_f32 is not "
+        "successful. Got status %d. Use -c dbg to see XNN logs.",
+        status);
+    return;
+  }
+
+  xnn_run_operator(softmax_op, nullptr /* thread pool */);
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
new file mode 100644
index 00000000000..52ac52a785b
--- /dev/null
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -0,0 +1,59 @@
+/* Copyright 2019 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.
+ * ===========================================================================*/
+
+#include <gtest/gtest.h>
+
+#include <cstddef>
+#include <vector>
+
+#include "src/cc/backend.h"
+#include "src/cc/kernels/Sigmoid.h"
+
+TEST(BATCH_MATMUL, xnn_operator_lifetime) {
+  tfjs::wasm::init();
+
+  ASSERT_EQ(0, tfjs::backend::num_tensors());
+
+  const size_t x0_id = 1;
+  const size_t x1_id = 2;
+  const size_t x_size = 4;
+  float x_values[x_size] = {1, 2, 2, 2};
+
+  const size_t out_id = 3;
+  const size_t out_size = 8;
+  float out_values[out_size] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+  tfjs::wasm::register_tensor(x0_id, x_size, x_values);
+  tfjs::wasm::register_tensor(x1_id, x_size, x_values);
+  tfjs::wasm::register_tensor(out_id, out_size, out_values);
+
+  ASSERT_EQ(3, tfjs::backend::num_tensors());
+  ASSERT_EQ(0, tfjs::backend::xnn_operator_count);
+
+  // One new xnn_operator should be created for the first call to Sigmoid.
+  tfjs::wasm::Sigmoid(x0_id, out_id);
+  ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
+
+  // No new xnn_operators should be created for the second call to
+  // Sigmoid with the same arguments.
+  tfjs::wasm::Sigmoid(x0_id, out_id);
+  ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
+
+  // No new xnn_operators should be created for the second call to
+  // Sigmoid with different arguments.
+  tfjs::wasm::Sigmoid(x1_id, out_id);
+  ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
+
+  tfjs::wasm::dispose();
+}
diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index 370615f2cf3..244bd8eccc2 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -19,6 +19,7 @@ import * as tf from '@tensorflow/tfjs-core';
 import {registerBackend, removeBackend, test_util} from '@tensorflow/tfjs-core';
 // tslint:disable-next-line:no-imports-from-dist
 import {ALL_ENVS, BROWSER_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
+import {expectArraysClose} from '@tensorflow/tfjs-core/dist/test_util';
 
 import {init, resetWasmPath} from './backend_wasm';
 import {BackendWasm, setWasmPath} from './index';
@@ -58,8 +59,8 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     }, 100);
 
     // Silences backend registration warnings.
-    spyOn(console, 'warn');
-    spyOn(console, 'log');
+    // spyOn(console, 'warn');
+    // spyOn(console, 'log');
   });
 
   afterEach(() => {
@@ -92,4 +93,14 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     expect(() => setWasmPath('too/late'))
         .toThrowError(/The WASM backend was already initialized. Make sure/);
   });
+
+  fit('softmax basic', async () => {
+    const y = tf.softmax(tf.tensor1d([2, 1, 3]));
+    const data = await y.data();
+    console.log('SOFTMAX BASIC');
+    console.log(Array.from(data));
+
+    expectArraysClose(data, [0.24472847, 0.09003057, 0.66524095]);
+    expectArraysClose(await y.sum().data(), 1);
+  });
 });
diff --git a/tfjs-backend-wasm/src/kernels/Softmax.ts b/tfjs-backend-wasm/src/kernels/Softmax.ts
new file mode 100644
index 00000000000..3e3e4a38949
--- /dev/null
+++ b/tfjs-backend-wasm/src/kernels/Softmax.ts
@@ -0,0 +1,54 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+
+import {BackendWasm} from '../backend_wasm';
+
+interface SoftmaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+let wasmFunc: (xId: number, outId: number) => void;
+
+function setup(backend: BackendWasm): void {
+  wasmFunc =
+      backend.wasm.cwrap('Softmax', null /* void */, ['number', 'number']);
+}
+
+function softmax(args: {backend: BackendWasm, inputs: SoftmaxInputs}):
+    TensorInfo {
+  const {backend, inputs: {logits}} = args;
+  const xId = backend.dataIdMap.get(logits.dataId).id;
+  const out = backend.makeOutput(logits.shape, logits.dtype);
+  const outId = backend.dataIdMap.get(out.dataId).id;
+
+  // Short-circuit zero-sized tensors.
+  if (util.sizeFromShape(out.shape) === 0) {
+    return out;
+  }
+
+  wasmFunc(xId, outId);
+  return out;
+}
+
+registerKernel({
+  kernelName: 'Softmax',
+  backendName: 'wasm',
+  setupFunc: setup,
+  kernelFunc: softmax
+});
diff --git a/tfjs-backend-wasm/src/kernels/all_kernels.ts b/tfjs-backend-wasm/src/kernels/all_kernels.ts
index 59de2028718..3b1fdd81d65 100644
--- a/tfjs-backend-wasm/src/kernels/all_kernels.ts
+++ b/tfjs-backend-wasm/src/kernels/all_kernels.ts
@@ -64,6 +64,7 @@ import './ScatterNd';
 import './Sigmoid';
 import './Sin';
 import './Slice';
+import './Softmax';
 import './Square';
 import './Sub';
 import './Sum';
diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts
index 9b794fa25f0..8c4d289b7e1 100644
--- a/tfjs-backend-wasm/src/setup_test.ts
+++ b/tfjs-backend-wasm/src/setup_test.ts
@@ -34,6 +34,7 @@ const TEST_FILTERS: TestFilter[] = [
       'Tensor2D float32 -> bool', 'Tensor2D int32 -> bool'
     ]
   },
+  {include: 'softmax'},
   {
     include: 'add ',
     excludes: [
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index b9144cf2af0..88902fc26f3 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -73,7 +73,7 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
           logits: () => dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y))
         };
       },
-      'Softmax', {dim});
+      'Softmax');
 }
 
 /**

From 5560bb815409329c9139cd4d3e1e4d17c806f7f0 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 14:14:10 -0500
Subject: [PATCH 08/84] wtf

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index 75c0b234ca8..60f58e73f55 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -60,8 +60,8 @@ void Softmax(size_t x_id, size_t out_id) {
 
   auto operator_cache_idx = operator_cache.find(cache_key);
   if (operator_cache_idx == operator_cache.end()) {
-    const size_t input_stride = 1;
-    const size_t output_stride = 1;
+    const size_t input_stride = channels;
+    const size_t output_stride = channels;
     const uint32_t flags = 0;
 
     xnn_status status = xnn_create_softmax_nc_f32(

From 8a59b0cdbf696fdae4d42a8a84f3833adb29b3b6 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 15:36:13 -0500
Subject: [PATCH 09/84] use

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index 60f58e73f55..080181e56e1 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -46,7 +46,7 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 
-void Softmax(size_t x_id, size_t out_id) {
+void Softmax(const size_t x_id, const size_t out_id) {
   auto& x_info = backend::get_tensor_info(x_id);
   auto& out_info = backend::get_tensor_info_out(out_id);
 
@@ -55,7 +55,7 @@ void Softmax(size_t x_id, size_t out_id) {
 
   xnn_operator_t softmax_op = nullptr;
 
-  const size_t channels = 3;
+  const size_t channels = x_info.size;
   OperatorCacheKey cache_key = {channels};
 
   auto operator_cache_idx = operator_cache.find(cache_key);
@@ -63,6 +63,7 @@ void Softmax(size_t x_id, size_t out_id) {
     const size_t input_stride = channels;
     const size_t output_stride = channels;
     const uint32_t flags = 0;
+    // input_stride / output_stride must be >= channels
 
     xnn_status status = xnn_create_softmax_nc_f32(
         channels, input_stride, output_stride, flags, &softmax_op);

From d5037f0bbe5b4d463759b6b5180052eba726a635 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 15:40:43 -0500
Subject: [PATCH 10/84] simplify

---
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 12 +++---------
 tfjs-core/src/backends/webgl/backend_webgl.ts | 12 +++---------
 2 files changed, 6 insertions(+), 18 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index 7f9dbd0e620..bc37e20c582 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -384,15 +384,9 @@ export class MathBackendCPU extends KernelBackend {
         logits,
         maxLogit.reshape(axis_util.expandShapeToKeepDim(maxLogit.shape, axes)));
     const b = this.exp(a);
-    const c = this.sum(b, axes);
-    const d = this.log(c);
-    let lse = this.add(this.reshape(maxLogit, d.shape), d);
-    const newShape = axis_util.expandShapeToKeepDim(lse.shape, axes);
-    lse = this.reshape(lse, newShape);
-
-    const logResult = this.subtract(logits, lse);
-    const y = this.exp(logResult);
-    return y as T;
+    const sumExp = this.sum(b, axes);
+
+    return b.div(sumExp);
   }
 
   subtract(a: Tensor, b: Tensor): Tensor {
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 1ca7e956e0a..343a3e1c496 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1595,15 +1595,9 @@ export class MathBackendWebGL extends KernelBackend {
         logits,
         maxLogit.reshape(axis_util.expandShapeToKeepDim(maxLogit.shape, axes)));
     const b = this.exp(a);
-    const c = this.sum(b, axes);
-    const d = this.log(c);
-    let lse = this.add(this.reshape(maxLogit, d.shape), d);
-    const newShape = axis_util.expandShapeToKeepDim(lse.shape, axes);
-    lse = this.reshape(lse, newShape);
-
-    const logResult = this.subtract(logits, lse);
-    const y = this.exp(logResult);
-    return y as T;
+    const sumExp = this.sum(b, axes);
+
+    return b.div(sumExp);
   }
 
   log<T extends Tensor>(x: T): T {

From 6981390bb56b9de01b65bbdaa274722de991d292 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 4 Feb 2020 16:28:12 -0500
Subject: [PATCH 11/84] softmax test

---
 tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
index 52ac52a785b..dd09153d4a0 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -18,9 +18,8 @@
 #include <vector>
 
 #include "src/cc/backend.h"
-#include "src/cc/kernels/Sigmoid.h"
 
-TEST(BATCH_MATMUL, xnn_operator_lifetime) {
+TEST(SOFTMAX, xnn_operator_lifetime) {
   tfjs::wasm::init();
 
   ASSERT_EQ(0, tfjs::backend::num_tensors());
@@ -41,18 +40,18 @@ TEST(BATCH_MATMUL, xnn_operator_lifetime) {
   ASSERT_EQ(3, tfjs::backend::num_tensors());
   ASSERT_EQ(0, tfjs::backend::xnn_operator_count);
 
-  // One new xnn_operator should be created for the first call to Sigmoid.
-  tfjs::wasm::Sigmoid(x0_id, out_id);
+  // One new xnn_operator should be created for the first call to Softmax.
+  tfjs::wasm::Softmax(x0_id, out_id);
   ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
 
   // No new xnn_operators should be created for the second call to
-  // Sigmoid with the same arguments.
-  tfjs::wasm::Sigmoid(x0_id, out_id);
+  // Softmax with the same arguments.
+  tfjs::wasm::Softmax(x0_id, out_id);
   ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
 
   // No new xnn_operators should be created for the second call to
-  // Sigmoid with different arguments.
-  tfjs::wasm::Sigmoid(x1_id, out_id);
+  // Softmax with different arguments.
+  tfjs::wasm::Softmax(x1_id, out_id);
   ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
 
   tfjs::wasm::dispose();

From 6b01d0ce1a5f108b51f6f817cd3c2bb48d78d209 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 08:21:41 -0500
Subject: [PATCH 12/84] fix

---
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 7 +++----
 tfjs-core/src/backends/webgl/backend_webgl.ts | 7 +++----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index bc37e20c582..099f9d5a973 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -380,11 +380,10 @@ export class MathBackendCPU extends KernelBackend {
   softmax<T extends Tensor>(logits: T): T {
     const axes = util.parseAxisParam([logits.rank - 1], logits.shape);
     const maxLogit = this.max(logits, axes);
-    const a = this.subtract(
-        logits,
-        maxLogit.reshape(axis_util.expandShapeToKeepDim(maxLogit.shape, axes)));
+    const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
+    const a = this.subtract(logits, maxLogit.reshape(expandedShape));
     const b = this.exp(a);
-    const sumExp = this.sum(b, axes);
+    const sumExp = this.sum(b, axes).reshape(expandedShape);
 
     return b.div(sumExp);
   }
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 343a3e1c496..39e23eea53d 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1591,11 +1591,10 @@ export class MathBackendWebGL extends KernelBackend {
   softmax<T extends Tensor>(logits: T): T {
     const axes = util.parseAxisParam([logits.rank - 1], logits.shape);
     const maxLogit = this.max(logits, axes);
-    const a = this.subtract(
-        logits,
-        maxLogit.reshape(axis_util.expandShapeToKeepDim(maxLogit.shape, axes)));
+    const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
+    const a = this.subtract(logits, maxLogit.reshape(expandedShape));
     const b = this.exp(a);
-    const sumExp = this.sum(b, axes);
+    const sumExp = this.sum(b, axes).reshape(expandedShape);
 
     return b.div(sumExp);
   }

From 099a594d3d1b89d79277e12811a97f15be21eb02 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 08:46:40 -0500
Subject: [PATCH 13/84] pass in dim

---
 tfjs-backend-wasm/src/Softmax_xnn.cc          | 100 ++++++++++++++++++
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc   |   8 +-
 tfjs-backend-wasm/src/kernels/Softmax.ts      |  27 +++--
 tfjs-core/src/backends/backend.ts             |   2 +-
 tfjs-core/src/backends/cpu/backend_cpu.ts     |   4 +-
 tfjs-core/src/backends/webgl/backend_webgl.ts |   4 +-
 tfjs-core/src/ops/softmax.ts                  |   4 +-
 7 files changed, 131 insertions(+), 18 deletions(-)
 create mode 100644 tfjs-backend-wasm/src/Softmax_xnn.cc

diff --git a/tfjs-backend-wasm/src/Softmax_xnn.cc b/tfjs-backend-wasm/src/Softmax_xnn.cc
new file mode 100644
index 00000000000..080181e56e1
--- /dev/null
+++ b/tfjs-backend-wasm/src/Softmax_xnn.cc
@@ -0,0 +1,100 @@
+/* Copyright 2020 Google LLC. 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <xnnpack.h>
+#include <cmath>
+#include <cstddef>
+#include <map>
+#include <tuple>
+#include <vector>
+
+#include "src/cc/backend.h"
+#include "src/cc/util.h"
+
+namespace {
+// We use std::tuple as the cache key as it implements the compare operator
+// needed for std::map.
+typedef std::tuple<size_t> OperatorCacheKey;
+
+// The operator cache maps the weights id to the xnn_operator_t instantiated for
+// this set of weights.
+std::map<OperatorCacheKey, xnn_operator_t> operator_cache;
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+
+void Softmax(const size_t x_id, const size_t out_id) {
+  auto& x_info = backend::get_tensor_info(x_id);
+  auto& out_info = backend::get_tensor_info_out(out_id);
+
+  const float* x_buf = x_info.f32();
+  float* out_buf = out_info.f32_write();
+
+  xnn_operator_t softmax_op = nullptr;
+
+  const size_t channels = x_info.size;
+  OperatorCacheKey cache_key = {channels};
+
+  auto operator_cache_idx = operator_cache.find(cache_key);
+  if (operator_cache_idx == operator_cache.end()) {
+    const size_t input_stride = channels;
+    const size_t output_stride = channels;
+    const uint32_t flags = 0;
+    // input_stride / output_stride must be >= channels
+
+    xnn_status status = xnn_create_softmax_nc_f32(
+        channels, input_stride, output_stride, flags, &softmax_op);
+    if (status != xnn_status_success) {
+      tfjs::util::warn(
+          "XNN status for xnn_create_softmax_nc_f32 is not "
+          "successful. Got status %d. Use -c dbg to see XNN logs.",
+          status);
+      return;
+    }
+
+    operator_cache.insert({cache_key, softmax_op});
+
+    tfjs::backend::xnn_operator_count++;
+  } else {
+    softmax_op = operator_cache_idx->second;
+  }
+
+  xnn_status status = xnn_setup_softmax_nc_f32(softmax_op, 1, x_buf, out_buf,
+                                               nullptr /* thread pool */);
+  if (status != xnn_status_success) {
+    tfjs::util::warn(
+        "XNN status for xnn_setup_softmax_nc_f32 is not "
+        "successful. Got status %d. Use -c dbg to see XNN logs.",
+        status);
+    return;
+  }
+
+  xnn_run_operator(softmax_op, nullptr /* thread pool */);
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index 080181e56e1..28cec79783d 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -46,7 +46,8 @@ extern "C" {
 EMSCRIPTEN_KEEPALIVE
 #endif
 
-void Softmax(const size_t x_id, const size_t out_id) {
+void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
+             const size_t batch) {
   auto& x_info = backend::get_tensor_info(x_id);
   auto& out_info = backend::get_tensor_info_out(out_id);
 
@@ -55,7 +56,6 @@ void Softmax(const size_t x_id, const size_t out_id) {
 
   xnn_operator_t softmax_op = nullptr;
 
-  const size_t channels = x_info.size;
   OperatorCacheKey cache_key = {channels};
 
   auto operator_cache_idx = operator_cache.find(cache_key);
@@ -82,8 +82,8 @@ void Softmax(const size_t x_id, const size_t out_id) {
     softmax_op = operator_cache_idx->second;
   }
 
-  xnn_status status = xnn_setup_softmax_nc_f32(softmax_op, 1, x_buf, out_buf,
-                                               nullptr /* thread pool */);
+  xnn_status status = xnn_setup_softmax_nc_f32(
+      softmax_op, batch, x_buf, out_buf, nullptr /* thread pool */);
   if (status != xnn_status_success) {
     tfjs::util::warn(
         "XNN status for xnn_setup_softmax_nc_f32 is not "
diff --git a/tfjs-backend-wasm/src/kernels/Softmax.ts b/tfjs-backend-wasm/src/kernels/Softmax.ts
index 3e3e4a38949..19eeee113f9 100644
--- a/tfjs-backend-wasm/src/kernels/Softmax.ts
+++ b/tfjs-backend-wasm/src/kernels/Softmax.ts
@@ -15,7 +15,7 @@
  * =============================================================================
  */
 
-import {NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
+import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo, util} from '@tensorflow/tfjs-core';
 
 import {BackendWasm} from '../backend_wasm';
 
@@ -23,26 +23,39 @@ interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
 }
 
-let wasmFunc: (xId: number, outId: number) => void;
+interface SoftmaxAttrs extends NamedAttrMap {
+  dim: number;
+}
+
+let wasmFunc: (xId: number, outId: number, channels: number, batch: number) =>
+    void;
 
 function setup(backend: BackendWasm): void {
-  wasmFunc =
-      backend.wasm.cwrap('Softmax', null /* void */, ['number', 'number']);
+  wasmFunc = backend.wasm.cwrap('Softmax', null /* void */, [
+    'number',  // xId
+    'number',  // outId
+    'number',  // channels
+    'number'   // batch
+  ]);
 }
 
-function softmax(args: {backend: BackendWasm, inputs: SoftmaxInputs}):
+function softmax(
+    args: {backend: BackendWasm, inputs: SoftmaxInputs, attrs: SoftmaxAttrs}):
     TensorInfo {
-  const {backend, inputs: {logits}} = args;
+  const {backend, inputs: {logits}, attrs: {dim}} = args;
   const xId = backend.dataIdMap.get(logits.dataId).id;
   const out = backend.makeOutput(logits.shape, logits.dtype);
   const outId = backend.dataIdMap.get(out.dataId).id;
 
+  const channels = logits.shape[dim];
+  const batch = util.sizeFromShape(logits.shape) / channels;
+
   // Short-circuit zero-sized tensors.
   if (util.sizeFromShape(out.shape) === 0) {
     return out;
   }
 
-  wasmFunc(xId, outId);
+  wasmFunc(xId, outId, channels, batch);
   return out;
 }
 
diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts
index 9901a3614db..5b92a1492a0 100644
--- a/tfjs-core/src/backends/backend.ts
+++ b/tfjs-core/src/backends/backend.ts
@@ -305,7 +305,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
   expm1<T extends Tensor>(x: T): T {
     return notYetImplemented('expm1');
   }
-  softmax<T extends Tensor>(x: T): T {
+  softmax<T extends Tensor>(x: T, dim: number): T {
     return notYetImplemented('softmax');
   }
   log<T extends Tensor>(x: T): T {
diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index 099f9d5a973..480770d17f4 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,8 +377,8 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
-  softmax<T extends Tensor>(logits: T): T {
-    const axes = util.parseAxisParam([logits.rank - 1], logits.shape);
+  softmax<T extends Tensor>(logits: T, dim: number): T {
+    const axes = util.parseAxisParam([dim], logits.shape);
     const maxLogit = this.max(logits, axes);
     const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
     const a = this.subtract(logits, maxLogit.reshape(expandedShape));
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 39e23eea53d..f7ffae96b2b 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1588,8 +1588,8 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  softmax<T extends Tensor>(logits: T): T {
-    const axes = util.parseAxisParam([logits.rank - 1], logits.shape);
+  softmax<T extends Tensor>(logits: T, dim: number): T {
+    const axes = util.parseAxisParam([dim], logits.shape);
     const maxLogit = this.max(logits, axes);
     const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
     const a = this.subtract(logits, maxLogit.reshape(expandedShape));
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index 88902fc26f3..b6933d97bd4 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -59,7 +59,7 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
 
   return ENGINE.runKernelFunc(
       (backend, save) => {
-        const y = backend.softmax($logits);
+        const y = backend.softmax($logits, dim);
         save([y]);
         return y;
       },
@@ -73,7 +73,7 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
           logits: () => dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y))
         };
       },
-      'Softmax');
+      'Softmax', {dim});
 }
 
 /**

From 188bde6866ea79a73ebb5fa3798ffb08386a853b Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 09:58:57 -0500
Subject: [PATCH 14/84] test case

---
 tfjs-backend-wasm/src/cc/BUILD                |  1 +
 tfjs-backend-wasm/src/cc/kernels/Softmax.h    | 31 +++++++++++++++++++
 .../src/cc/kernels/Softmax_test.cc            | 21 ++++++++++---
 tfjs-backend-wasm/src/index_test.ts           |  7 ++---
 4 files changed, 51 insertions(+), 9 deletions(-)
 create mode 100644 tfjs-backend-wasm/src/cc/kernels/Softmax.h

diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index 4209ce300b6..30adcff5e1a 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -651,6 +651,7 @@ tfjs_cc_library(
 tfjs_cc_library(
     name = "Softmax",
     srcs = ["kernels/Softmax.cc"],
+    hdrs = ["kernels/Softmax.h"],
     deps = [
         ":backend",
         ":unary",
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.h b/tfjs-backend-wasm/src/cc/kernels/Softmax.h
new file mode 100644
index 00000000000..2f6ab326c21
--- /dev/null
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.h
@@ -0,0 +1,31 @@
+/* Copyright 2019 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.
+ * ===========================================================================*/
+
+#ifndef KERNELS_SOFTMAX_H_
+#define KERNELS_SOFTMAX_H_
+
+#include <cstddef>
+
+namespace tfjs {
+namespace wasm {
+extern "C" {
+
+void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
+             const size_t batch);
+}
+
+}  // namespace wasm
+}  // namespace tfjs
+
+#endif  // KERNELS_SOFTMAX_H_
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
index dd09153d4a0..f64f3506384 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -18,6 +18,8 @@
 #include <vector>
 
 #include "src/cc/backend.h"
+#include "src/cc/kernels/Softmax.h"
+#include "src/cc/util.h"
 
 TEST(SOFTMAX, xnn_operator_lifetime) {
   tfjs::wasm::init();
@@ -30,8 +32,8 @@ TEST(SOFTMAX, xnn_operator_lifetime) {
   float x_values[x_size] = {1, 2, 2, 2};
 
   const size_t out_id = 3;
-  const size_t out_size = 8;
-  float out_values[out_size] = {0, 0, 0, 0, 0, 0, 0, 0};
+  const size_t out_size = 4;
+  float out_values[out_size] = {0, 0, 0, 0};
 
   tfjs::wasm::register_tensor(x0_id, x_size, x_values);
   tfjs::wasm::register_tensor(x1_id, x_size, x_values);
@@ -41,17 +43,26 @@ TEST(SOFTMAX, xnn_operator_lifetime) {
   ASSERT_EQ(0, tfjs::backend::xnn_operator_count);
 
   // One new xnn_operator should be created for the first call to Softmax.
-  tfjs::wasm::Softmax(x0_id, out_id);
+  tfjs::wasm::Softmax(x0_id, out_id, 4, 1);
   ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
 
+  auto& out_info = tfjs::backend::get_tensor_info(out_id);
+  const float* out_buf = out_info.f32();
+
+  tfjs::util::log("PRINTING OUT VALUES");
+  for (size_t i = 0; i < out_size; ++i) {
+    tfjs::util::log("%f", out_buf[i]);
+  }
+  // prints: 0.109232, 0.296923, 0.296923, 0.296923
+
   // No new xnn_operators should be created for the second call to
   // Softmax with the same arguments.
-  tfjs::wasm::Softmax(x0_id, out_id);
+  tfjs::wasm::Softmax(x0_id, out_id, 4, 1);
   ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
 
   // No new xnn_operators should be created for the second call to
   // Softmax with different arguments.
-  tfjs::wasm::Softmax(x1_id, out_id);
+  tfjs::wasm::Softmax(x1_id, out_id, 4, 1);
   ASSERT_EQ(1, tfjs::backend::xnn_operator_count);
 
   tfjs::wasm::dispose();
diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index 244bd8eccc2..292afe047b5 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -95,12 +95,11 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
   });
 
   fit('softmax basic', async () => {
-    const y = tf.softmax(tf.tensor1d([2, 1, 3]));
+    const y = tf.softmax(tf.tensor1d([1, 2, 2, 2]));
     const data = await y.data();
-    console.log('SOFTMAX BASIC');
-    console.log(Array.from(data));
+    console.log(Array.from(data));  // [0.1428571, 0.285714, 0.285714, 0.285714]
 
-    expectArraysClose(data, [0.24472847, 0.09003057, 0.66524095]);
+    expectArraysClose(data, [0.1092318, 0.2969227, 0.2969227, 0.2969227]);
     expectArraysClose(await y.sum().data(), 1);
   });
 });

From 19a91bdb272fa73d6aefe6f9bea738c7c7b48f33 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 10:03:37 -0500
Subject: [PATCH 15/84] logs

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index 28cec79783d..f5b5301de01 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -54,6 +54,14 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
   const float* x_buf = x_info.f32();
   float* out_buf = out_info.f32_write();
 
+  tfjs::util::log("XNN INPUT");
+  for (size_t i = 0; i < batch; ++i) {
+    for (size_t j = 0; j < channels; ++j) {
+      tfjs::util::log("%f", x_buf[i * channels + j]);
+    }
+  }
+  // prints: 1.0, 2.0, 2.0, 2.0
+
   xnn_operator_t softmax_op = nullptr;
 
   OperatorCacheKey cache_key = {channels};
@@ -93,6 +101,14 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
   }
 
   xnn_run_operator(softmax_op, nullptr /* thread pool */);
+
+  tfjs::util::log("XNN OUTPUT");
+  for (size_t i = 0; i < batch; ++i) {
+    for (size_t j = 0; j < channels; ++j) {
+      tfjs::util::log("%f", out_buf[i * channels + j]);
+    }
+  }
+  // prints: 0.1428571, 0.285714, 0.285714, 0.285714
 }
 
 }  // extern "C"

From 74af25163b4608df057dcbaff981af1f8efd969f Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 10:07:17 -0500
Subject: [PATCH 16/84] delete

---
 tfjs-backend-wasm/src/Softmax_xnn.cc | 100 ---------------------------
 1 file changed, 100 deletions(-)
 delete mode 100644 tfjs-backend-wasm/src/Softmax_xnn.cc

diff --git a/tfjs-backend-wasm/src/Softmax_xnn.cc b/tfjs-backend-wasm/src/Softmax_xnn.cc
deleted file mode 100644
index 080181e56e1..00000000000
--- a/tfjs-backend-wasm/src/Softmax_xnn.cc
+++ /dev/null
@@ -1,100 +0,0 @@
-/* Copyright 2020 Google LLC. 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.
- * ===========================================================================*/
-
-#ifdef __EMSCRIPTEN__
-#include <emscripten.h>
-#endif
-
-#include <xnnpack.h>
-#include <cmath>
-#include <cstddef>
-#include <map>
-#include <tuple>
-#include <vector>
-
-#include "src/cc/backend.h"
-#include "src/cc/util.h"
-
-namespace {
-// We use std::tuple as the cache key as it implements the compare operator
-// needed for std::map.
-typedef std::tuple<size_t> OperatorCacheKey;
-
-// The operator cache maps the weights id to the xnn_operator_t instantiated for
-// this set of weights.
-std::map<OperatorCacheKey, xnn_operator_t> operator_cache;
-
-}  // namespace
-
-namespace tfjs {
-namespace wasm {
-
-extern "C" {
-
-#ifdef __EMSCRIPTEN__
-EMSCRIPTEN_KEEPALIVE
-#endif
-
-void Softmax(const size_t x_id, const size_t out_id) {
-  auto& x_info = backend::get_tensor_info(x_id);
-  auto& out_info = backend::get_tensor_info_out(out_id);
-
-  const float* x_buf = x_info.f32();
-  float* out_buf = out_info.f32_write();
-
-  xnn_operator_t softmax_op = nullptr;
-
-  const size_t channels = x_info.size;
-  OperatorCacheKey cache_key = {channels};
-
-  auto operator_cache_idx = operator_cache.find(cache_key);
-  if (operator_cache_idx == operator_cache.end()) {
-    const size_t input_stride = channels;
-    const size_t output_stride = channels;
-    const uint32_t flags = 0;
-    // input_stride / output_stride must be >= channels
-
-    xnn_status status = xnn_create_softmax_nc_f32(
-        channels, input_stride, output_stride, flags, &softmax_op);
-    if (status != xnn_status_success) {
-      tfjs::util::warn(
-          "XNN status for xnn_create_softmax_nc_f32 is not "
-          "successful. Got status %d. Use -c dbg to see XNN logs.",
-          status);
-      return;
-    }
-
-    operator_cache.insert({cache_key, softmax_op});
-
-    tfjs::backend::xnn_operator_count++;
-  } else {
-    softmax_op = operator_cache_idx->second;
-  }
-
-  xnn_status status = xnn_setup_softmax_nc_f32(softmax_op, 1, x_buf, out_buf,
-                                               nullptr /* thread pool */);
-  if (status != xnn_status_success) {
-    tfjs::util::warn(
-        "XNN status for xnn_setup_softmax_nc_f32 is not "
-        "successful. Got status %d. Use -c dbg to see XNN logs.",
-        status);
-    return;
-  }
-
-  xnn_run_operator(softmax_op, nullptr /* thread pool */);
-}
-
-}  // extern "C"
-}  // namespace wasm
-}  // namespace tfjs

From 007edeba41c0432894f97caee35e20c76d09ebd8 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 10:16:41 -0500
Subject: [PATCH 17/84] add batch to key

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index f5b5301de01..e43bc9e795f 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -29,7 +29,7 @@
 namespace {
 // We use std::tuple as the cache key as it implements the compare operator
 // needed for std::map.
-typedef std::tuple<size_t> OperatorCacheKey;
+typedef std::tuple<size_t, size_t> OperatorCacheKey;
 
 // The operator cache maps the weights id to the xnn_operator_t instantiated for
 // this set of weights.
@@ -64,14 +64,13 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
 
   xnn_operator_t softmax_op = nullptr;
 
-  OperatorCacheKey cache_key = {channels};
+  OperatorCacheKey cache_key = {channels, batch};
 
   auto operator_cache_idx = operator_cache.find(cache_key);
   if (operator_cache_idx == operator_cache.end()) {
     const size_t input_stride = channels;
     const size_t output_stride = channels;
     const uint32_t flags = 0;
-    // input_stride / output_stride must be >= channels
 
     xnn_status status = xnn_create_softmax_nc_f32(
         channels, input_stride, output_stride, flags, &softmax_op);

From bbcea3074d83d7e10ba0481e9431b372ba56b9d9 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 5 Feb 2020 18:13:16 -0500
Subject: [PATCH 18/84] move log statement

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index e43bc9e795f..0c380732069 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -54,14 +54,6 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
   const float* x_buf = x_info.f32();
   float* out_buf = out_info.f32_write();
 
-  tfjs::util::log("XNN INPUT");
-  for (size_t i = 0; i < batch; ++i) {
-    for (size_t j = 0; j < channels; ++j) {
-      tfjs::util::log("%f", x_buf[i * channels + j]);
-    }
-  }
-  // prints: 1.0, 2.0, 2.0, 2.0
-
   xnn_operator_t softmax_op = nullptr;
 
   OperatorCacheKey cache_key = {channels, batch};
@@ -99,6 +91,14 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
     return;
   }
 
+  tfjs::util::log("XNN INPUT");
+  for (size_t i = 0; i < batch; ++i) {
+    for (size_t j = 0; j < channels; ++j) {
+      tfjs::util::log("%f", x_buf[i * channels + j]);
+    }
+  }
+  // prints: 1.0, 2.0, 2.0, 2.0
+
   xnn_run_operator(softmax_op, nullptr /* thread pool */);
 
   tfjs::util::log("XNN OUTPUT");

From 932e7a69b9a09fbef992ae44765d85fdeb6681bc Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 6 Feb 2020 10:54:21 -0500
Subject: [PATCH 19/84] remove batch from cache

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index 0c380732069..3905ba3a9d6 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -29,7 +29,7 @@
 namespace {
 // We use std::tuple as the cache key as it implements the compare operator
 // needed for std::map.
-typedef std::tuple<size_t, size_t> OperatorCacheKey;
+typedef std::tuple<size_t> OperatorCacheKey;
 
 // The operator cache maps the weights id to the xnn_operator_t instantiated for
 // this set of weights.
@@ -56,7 +56,7 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
 
   xnn_operator_t softmax_op = nullptr;
 
-  OperatorCacheKey cache_key = {channels, batch};
+  OperatorCacheKey cache_key = {channels};
 
   auto operator_cache_idx = operator_cache.find(cache_key);
   if (operator_cache_idx == operator_cache.end()) {
@@ -99,9 +99,16 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
   }
   // prints: 1.0, 2.0, 2.0, 2.0
 
+  tfjs::util::log("XNN OUTPUT BEFORE");
+  for (size_t i = 0; i < batch; ++i) {
+    for (size_t j = 0; j < channels; ++j) {
+      tfjs::util::log("%f", out_buf[i * channels + j]);
+    }
+  }
+
   xnn_run_operator(softmax_op, nullptr /* thread pool */);
 
-  tfjs::util::log("XNN OUTPUT");
+  tfjs::util::log("XNN OUTPUT AFTER");
   for (size_t i = 0; i < batch; ++i) {
     for (size_t j = 0; j < channels; ++j) {
       tfjs::util::log("%f", out_buf[i * channels + j]);

From ccebc1ef37896f7472cee1a8099b84476a41a388 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 6 Feb 2020 11:55:01 -0500
Subject: [PATCH 20/84] add note

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index 3905ba3a9d6..dba78acb92b 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -105,6 +105,7 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
       tfjs::util::log("%f", out_buf[i * channels + j]);
     }
   }
+  // prints: 0, 0, 0, 0
 
   xnn_run_operator(softmax_op, nullptr /* thread pool */);
 

From fcb079ab0a8fd817a00ae8a2a06217c284667918 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 7 Feb 2020 13:02:46 -0500
Subject: [PATCH 21/84] remove build flags

---
 tfjs-backend-wasm/.bazelrc | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tfjs-backend-wasm/.bazelrc b/tfjs-backend-wasm/.bazelrc
index 0754893952d..8d066759084 100644
--- a/tfjs-backend-wasm/.bazelrc
+++ b/tfjs-backend-wasm/.bazelrc
@@ -13,8 +13,6 @@ build:wasm --cxxopt="-std=c++11"
 build:wasm --cxxopt="-fno-rtti"
 build:wasm --cxxopt="-fno-exceptions"
 build:wasm --cxxopt="-fomit-frame-pointer"
-build:wasm --cxxopt="-ffast-math"
-build:wasm --copt="-ffast-math"
 
 # Disable sandbox environment because emsdk caches files by writing to
 # home directory.

From 2daaa4ffc704184999bdeaa2de0354e03579d955 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 7 Feb 2020 17:35:25 -0500
Subject: [PATCH 22/84] remove logs

---
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc   | 24 -------------------
 .../src/cc/kernels/Softmax_test.cc            |  6 -----
 2 files changed, 30 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index dba78acb92b..c1cbe6ad3b6 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -91,31 +91,7 @@ void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
     return;
   }
 
-  tfjs::util::log("XNN INPUT");
-  for (size_t i = 0; i < batch; ++i) {
-    for (size_t j = 0; j < channels; ++j) {
-      tfjs::util::log("%f", x_buf[i * channels + j]);
-    }
-  }
-  // prints: 1.0, 2.0, 2.0, 2.0
-
-  tfjs::util::log("XNN OUTPUT BEFORE");
-  for (size_t i = 0; i < batch; ++i) {
-    for (size_t j = 0; j < channels; ++j) {
-      tfjs::util::log("%f", out_buf[i * channels + j]);
-    }
-  }
-  // prints: 0, 0, 0, 0
-
   xnn_run_operator(softmax_op, nullptr /* thread pool */);
-
-  tfjs::util::log("XNN OUTPUT AFTER");
-  for (size_t i = 0; i < batch; ++i) {
-    for (size_t j = 0; j < channels; ++j) {
-      tfjs::util::log("%f", out_buf[i * channels + j]);
-    }
-  }
-  // prints: 0.1428571, 0.285714, 0.285714, 0.285714
 }
 
 }  // extern "C"
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
index f64f3506384..2f59b8cfc0d 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -49,12 +49,6 @@ TEST(SOFTMAX, xnn_operator_lifetime) {
   auto& out_info = tfjs::backend::get_tensor_info(out_id);
   const float* out_buf = out_info.f32();
 
-  tfjs::util::log("PRINTING OUT VALUES");
-  for (size_t i = 0; i < out_size; ++i) {
-    tfjs::util::log("%f", out_buf[i]);
-  }
-  // prints: 0.109232, 0.296923, 0.296923, 0.296923
-
   // No new xnn_operators should be created for the second call to
   // Softmax with the same arguments.
   tfjs::wasm::Softmax(x0_id, out_id, 4, 1);

From dd0ba9f765f95c6723fe13519b4dbc8a2715f469 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 7 Feb 2020 17:39:51 -0500
Subject: [PATCH 23/84] testing

---
 tfjs-backend-wasm/src/index_test.ts | 2 +-
 tfjs-backend-wasm/src/setup_test.ts | 7 ++++++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index 292afe047b5..8bb185e0f1d 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -94,7 +94,7 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
         .toThrowError(/The WASM backend was already initialized. Make sure/);
   });
 
-  fit('softmax basic', async () => {
+  it('softmax basic', async () => {
     const y = tf.softmax(tf.tensor1d([1, 2, 2, 2]));
     const data = await y.data();
     console.log(Array.from(data));  // [0.1428571, 0.285714, 0.285714, 0.285714]
diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts
index 8c4d289b7e1..f1a27e41dc4 100644
--- a/tfjs-backend-wasm/src/setup_test.ts
+++ b/tfjs-backend-wasm/src/setup_test.ts
@@ -34,7 +34,12 @@ const TEST_FILTERS: TestFilter[] = [
       'Tensor2D float32 -> bool', 'Tensor2D int32 -> bool'
     ]
   },
-  {include: 'softmax'},
+  {
+    include: 'softmax',
+    excludes: [
+      'gradient'  // Gradient not yet implemented.
+    ]
+  },
   {
     include: 'add ',
     excludes: [

From 7962041d289f06676fd0183482518bdbfffcb4c4 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 7 Feb 2020 17:51:28 -0500
Subject: [PATCH 24/84] remove header

---
 tfjs-backend-wasm/src/cc/BUILD                |  1 -
 tfjs-backend-wasm/src/cc/kernels/Softmax.h    | 31 -------------------
 .../src/cc/kernels/Softmax_test.cc            |  1 -
 3 files changed, 33 deletions(-)
 delete mode 100644 tfjs-backend-wasm/src/cc/kernels/Softmax.h

diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index 30adcff5e1a..4209ce300b6 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -651,7 +651,6 @@ tfjs_cc_library(
 tfjs_cc_library(
     name = "Softmax",
     srcs = ["kernels/Softmax.cc"],
-    hdrs = ["kernels/Softmax.h"],
     deps = [
         ":backend",
         ":unary",
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.h b/tfjs-backend-wasm/src/cc/kernels/Softmax.h
deleted file mode 100644
index 2f6ab326c21..00000000000
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.h
+++ /dev/null
@@ -1,31 +0,0 @@
-/* Copyright 2019 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.
- * ===========================================================================*/
-
-#ifndef KERNELS_SOFTMAX_H_
-#define KERNELS_SOFTMAX_H_
-
-#include <cstddef>
-
-namespace tfjs {
-namespace wasm {
-extern "C" {
-
-void Softmax(const size_t x_id, const size_t out_id, const size_t channels,
-             const size_t batch);
-}
-
-}  // namespace wasm
-}  // namespace tfjs
-
-#endif  // KERNELS_SOFTMAX_H_
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
index 2f59b8cfc0d..e010471bdf2 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -18,7 +18,6 @@
 #include <vector>
 
 #include "src/cc/backend.h"
-#include "src/cc/kernels/Softmax.h"
 #include "src/cc/util.h"
 
 TEST(SOFTMAX, xnn_operator_lifetime) {

From 4e7fdc8ed54d12a4d96f0c73f27b1d599c92ca7d Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 7 Feb 2020 18:22:09 -0500
Subject: [PATCH 25/84] add neg

---
 tfjs-backend-wasm/src/cc/BUILD               |  9 +++++
 tfjs-backend-wasm/src/cc/kernels/Neg.cc      | 40 ++++++++++++++++++++
 tfjs-backend-wasm/src/kernels/Neg.ts         | 19 ++++++++++
 tfjs-backend-wasm/src/kernels/all_kernels.ts |  1 +
 tfjs-core/src/ops/unary_ops.ts               |  6 ++-
 5 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 tfjs-backend-wasm/src/cc/kernels/Neg.cc
 create mode 100644 tfjs-backend-wasm/src/kernels/Neg.ts

diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index 4209ce300b6..40a9836a619 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -543,6 +543,15 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "Neg",
+    srcs = ["kernels/Neg.cc"],
+    deps = [
+        ":backend",
+        ":unary",
+    ],
+)
+
 tfjs_cc_library(
     name = "NonMaxSuppressionV3",
     srcs = ["kernels/NonMaxSuppressionV3.cc"],
diff --git a/tfjs-backend-wasm/src/cc/kernels/Neg.cc b/tfjs-backend-wasm/src/cc/kernels/Neg.cc
new file mode 100644
index 00000000000..7570e1cc810
--- /dev/null
+++ b/tfjs-backend-wasm/src/cc/kernels/Neg.cc
@@ -0,0 +1,40 @@
+/* Copyright 2020 Google LLC. 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <cmath>
+
+#include "src/cc/backend.h"
+#include "src/cc/unary.h"
+
+namespace {
+inline float neg(const float val) { return -val; }
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void Neg(const int x_id, const int out_id) { unary(x_id, out_id, neg); }
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs
diff --git a/tfjs-backend-wasm/src/kernels/Neg.ts b/tfjs-backend-wasm/src/kernels/Neg.ts
new file mode 100644
index 00000000000..b226e799331
--- /dev/null
+++ b/tfjs-backend-wasm/src/kernels/Neg.ts
@@ -0,0 +1,19 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {registerUnaryKernel} from './unary_kernel';
+registerUnaryKernel('Neg');
diff --git a/tfjs-backend-wasm/src/kernels/all_kernels.ts b/tfjs-backend-wasm/src/kernels/all_kernels.ts
index 3b1fdd81d65..10be92b4869 100644
--- a/tfjs-backend-wasm/src/kernels/all_kernels.ts
+++ b/tfjs-backend-wasm/src/kernels/all_kernels.ts
@@ -51,6 +51,7 @@ import './MaxPool';
 import './Min';
 import './Minimum';
 import './Mul';
+import './Neg';
 import './NonMaxSuppressionV3';
 import './NonMaxSuppressionV5';
 import './PadV2';
diff --git a/tfjs-core/src/ops/unary_ops.ts b/tfjs-core/src/ops/unary_ops.ts
index beeb51f600f..3b13f7f53a1 100644
--- a/tfjs-core/src/ops/unary_ops.ts
+++ b/tfjs-core/src/ops/unary_ops.ts
@@ -42,7 +42,11 @@ function neg_<T extends Tensor>(x: T|TensorLike): T {
   const grad = (dy: T) => {
     return {$x: () => dy.neg()};
   };
-  return ENGINE.runKernelFunc(backend => backend.neg($x), {$x}, grad);
+
+  const attrs = {};
+  const inputsToSave = [$x];
+  return ENGINE.runKernelFunc(
+      backend => backend.neg($x), {$x}, grad, 'Neg', attrs, inputsToSave);
 }
 
 /**

From 869768b82044338c2a92d153997fcd5da10d00ca Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 7 Feb 2020 19:28:10 -0500
Subject: [PATCH 26/84] register

---
 tfjs-core/src/ops/unary_ops.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tfjs-core/src/ops/unary_ops.ts b/tfjs-core/src/ops/unary_ops.ts
index 3b13f7f53a1..7a33f1e0b42 100644
--- a/tfjs-core/src/ops/unary_ops.ts
+++ b/tfjs-core/src/ops/unary_ops.ts
@@ -40,13 +40,13 @@ function neg_<T extends Tensor>(x: T|TensorLike): T {
   const $x = convertToTensor(x, 'x', 'neg');
 
   const grad = (dy: T) => {
-    return {$x: () => dy.neg()};
+    return {x: () => dy.neg()};
   };
 
   const attrs = {};
   const inputsToSave = [$x];
   return ENGINE.runKernelFunc(
-      backend => backend.neg($x), {$x}, grad, 'Neg', attrs, inputsToSave);
+      backend => backend.neg($x), {x: $x}, grad, 'Neg', attrs, inputsToSave);
 }
 
 /**

From 3c9dda642269c1f8dcd30e30962d4193dc296250 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Sat, 8 Feb 2020 09:10:37 -0500
Subject: [PATCH 27/84] notequal

---
 tfjs-backend-wasm/src/cc/BUILD               | 11 ++++
 tfjs-backend-wasm/src/cc/kernels/NotEqual.cc | 60 ++++++++++++++++++++
 tfjs-backend-wasm/src/kernels/NotEqual.ts    | 20 +++++++
 tfjs-backend-wasm/src/kernels/all_kernels.ts |  1 +
 tfjs-core/src/ops/compare.ts                 |  5 +-
 5 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 tfjs-backend-wasm/src/cc/kernels/NotEqual.cc
 create mode 100644 tfjs-backend-wasm/src/kernels/NotEqual.ts

diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index 40a9836a619..62c4b2bd315 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -177,6 +177,7 @@ tfjs_cc_library(
         ":Min",
         ":Minimum",
         ":Mul",
+        ":NotEqual",
         ":NonMaxSuppressionV3",
         ":NonMaxSuppressionV5",
         ":PadV2",
@@ -458,6 +459,16 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "NotEqual",
+    srcs = ["kernels/NotEqual.cc"],
+    deps = [
+        ":backend",
+        ":binary",
+        ":util",
+    ],
+)
+
 tfjs_cc_library(
     name = "LogicalAnd",
     srcs = ["kernels/LogicalAnd.cc"],
diff --git a/tfjs-backend-wasm/src/cc/kernels/NotEqual.cc b/tfjs-backend-wasm/src/cc/kernels/NotEqual.cc
new file mode 100644
index 00000000000..064bb4e6194
--- /dev/null
+++ b/tfjs-backend-wasm/src/cc/kernels/NotEqual.cc
@@ -0,0 +1,60 @@
+/* Copyright 2020 Google LLC. 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.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include "src/cc/binary.h"
+#include "src/cc/util.h"
+
+namespace {
+template <class T>
+inline bool notEqual(T a, T b) {
+  return a != b;
+}
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void NotEqual(const int a_id, const size_t* a_shape_ptr, const int a_shape_len,
+              const int b_id, const size_t* b_shape_ptr, const int b_shape_len,
+              const DType input_type, const int out_id) {
+  switch (input_type) {
+    case DType::float32:
+      compare_f32(a_id, b_id, out_id, notEqual<float>);
+      break;
+    case DType::int32:
+      compare_i32(a_id, b_id, out_id, notEqual<int>);
+      break;
+    case DType::boolean:
+      compare_bool(a_id, b_id, out_id, notEqual<bool>);
+      break;
+    default:
+      util::warn(
+          "NotEqual for tensor ids %d and %d failed."
+          "Unsupported input_type %d",
+          a_id, b_id, input_type);
+  }
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs
diff --git a/tfjs-backend-wasm/src/kernels/NotEqual.ts b/tfjs-backend-wasm/src/kernels/NotEqual.ts
new file mode 100644
index 00000000000..a20be2a92d8
--- /dev/null
+++ b/tfjs-backend-wasm/src/kernels/NotEqual.ts
@@ -0,0 +1,20 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {registerBinaryKernel} from './binary_kernel';
+const supportsFullBroadcast = false;
+registerBinaryKernel('NotEqual', supportsFullBroadcast, 'bool');
diff --git a/tfjs-backend-wasm/src/kernels/all_kernels.ts b/tfjs-backend-wasm/src/kernels/all_kernels.ts
index 10be92b4869..157a4cbd8fa 100644
--- a/tfjs-backend-wasm/src/kernels/all_kernels.ts
+++ b/tfjs-backend-wasm/src/kernels/all_kernels.ts
@@ -54,6 +54,7 @@ import './Mul';
 import './Neg';
 import './NonMaxSuppressionV3';
 import './NonMaxSuppressionV5';
+import './NotEqual';
 import './PadV2';
 import './Prelu';
 import './Relu';
diff --git a/tfjs-core/src/ops/compare.ts b/tfjs-core/src/ops/compare.ts
index 7b32b7d8b89..521d269a8ef 100644
--- a/tfjs-core/src/ops/compare.ts
+++ b/tfjs-core/src/ops/compare.ts
@@ -47,8 +47,9 @@ function notEqual_<T extends Tensor>(
   let $b = convertToTensor(b, 'b', 'notEqual');
   [$a, $b] = makeTypesMatch($a, $b);
   assertAndGetBroadcastShape($a.shape, $b.shape);
-  return ENGINE.runKernelFunc(backend => backend.notEqual($a, $b), {$a, $b}) as
-      T;
+  return ENGINE.runKernelFunc(
+             backend => backend.notEqual($a, $b), {a: $a, b: $b},
+             null /* grad */, 'NotEqual') as T;
 }
 
 /**

From 9c58b334200e91653f599402959b84e16d7088cd Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Mon, 10 Feb 2020 10:44:35 -0500
Subject: [PATCH 28/84] lint

---
 tfjs-backend-wasm/src/index_test.ts | 10 ----------
 tfjs-core/src/ops/softmax.ts        |  1 -
 2 files changed, 11 deletions(-)

diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index 8bb185e0f1d..af58f387b71 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -19,7 +19,6 @@ import * as tf from '@tensorflow/tfjs-core';
 import {registerBackend, removeBackend, test_util} from '@tensorflow/tfjs-core';
 // tslint:disable-next-line:no-imports-from-dist
 import {ALL_ENVS, BROWSER_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
-import {expectArraysClose} from '@tensorflow/tfjs-core/dist/test_util';
 
 import {init, resetWasmPath} from './backend_wasm';
 import {BackendWasm, setWasmPath} from './index';
@@ -93,13 +92,4 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     expect(() => setWasmPath('too/late'))
         .toThrowError(/The WASM backend was already initialized. Make sure/);
   });
-
-  it('softmax basic', async () => {
-    const y = tf.softmax(tf.tensor1d([1, 2, 2, 2]));
-    const data = await y.data();
-    console.log(Array.from(data));  // [0.1428571, 0.285714, 0.285714, 0.285714]
-
-    expectArraysClose(data, [0.1092318, 0.2969227, 0.2969227, 0.2969227]);
-    expectArraysClose(await y.sum().data(), 1);
-  });
 });
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index b6933d97bd4..cfc837e0e74 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -24,7 +24,6 @@ import {TensorLike} from '../types';
 
 import {op} from './operation';
 
-
 /**
  * Computes the softmax normalized vector given the logits.
  *

From 88d8705e07fa5b140c1b0296a58664029e878fe3 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Mon, 10 Feb 2020 10:54:16 -0500
Subject: [PATCH 29/84] revive spy

---
 tfjs-backend-wasm/src/index_test.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tfjs-backend-wasm/src/index_test.ts b/tfjs-backend-wasm/src/index_test.ts
index af58f387b71..370615f2cf3 100644
--- a/tfjs-backend-wasm/src/index_test.ts
+++ b/tfjs-backend-wasm/src/index_test.ts
@@ -58,8 +58,8 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
     }, 100);
 
     // Silences backend registration warnings.
-    // spyOn(console, 'warn');
-    // spyOn(console, 'log');
+    spyOn(console, 'warn');
+    spyOn(console, 'log');
   });
 
   afterEach(() => {

From 648bc6c2003ea7a80b83d8db61da8efdbaacc6c4 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Mon, 10 Feb 2020 11:17:37 -0500
Subject: [PATCH 30/84] save


From 40d5e820c29fd267223c59657fd897408f6017ab Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Mon, 10 Feb 2020 12:36:54 -0500
Subject: [PATCH 31/84] add neg

---
 tfjs-backend-wasm/src/cc/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index 62c4b2bd315..e8e74ce1e52 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -177,6 +177,7 @@ tfjs_cc_library(
         ":Min",
         ":Minimum",
         ":Mul",
+        ":Neg",
         ":NotEqual",
         ":NonMaxSuppressionV3",
         ":NonMaxSuppressionV5",

From e149666c05797a515253b220a6efb94e6a210763 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 11 Feb 2020 10:08:08 -0500
Subject: [PATCH 32/84] save outputs

---
 tfjs-backend-wasm/src/setup_test.ts | 7 +------
 tfjs-core/src/ops/softmax.ts        | 5 ++++-
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/tfjs-backend-wasm/src/setup_test.ts b/tfjs-backend-wasm/src/setup_test.ts
index f1a27e41dc4..8c4d289b7e1 100644
--- a/tfjs-backend-wasm/src/setup_test.ts
+++ b/tfjs-backend-wasm/src/setup_test.ts
@@ -34,12 +34,7 @@ const TEST_FILTERS: TestFilter[] = [
       'Tensor2D float32 -> bool', 'Tensor2D int32 -> bool'
     ]
   },
-  {
-    include: 'softmax',
-    excludes: [
-      'gradient'  // Gradient not yet implemented.
-    ]
-  },
+  {include: 'softmax'},
   {
     include: 'add ',
     excludes: [
diff --git a/tfjs-core/src/ops/softmax.ts b/tfjs-core/src/ops/softmax.ts
index cfc837e0e74..fff8a3df23c 100644
--- a/tfjs-core/src/ops/softmax.ts
+++ b/tfjs-core/src/ops/softmax.ts
@@ -56,6 +56,9 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
         `Logits was rank ${$logits.rank} and dim was ${dim}`);
   }
 
+  const inputsToSave: Tensor[] = [];
+  const outputsToSave = [true];
+
   return ENGINE.runKernelFunc(
       (backend, save) => {
         const y = backend.softmax($logits, dim);
@@ -72,7 +75,7 @@ function softmax_<T extends Tensor>(logits: T|TensorLike, dim = -1): T {
           logits: () => dyTimesY.sub(dyTimesY.sum([dim], keepDims).mul(y))
         };
       },
-      'Softmax', {dim});
+      'Softmax', {dim}, inputsToSave, outputsToSave);
 }
 
 /**

From 8903f08e6a95692ea7ae6fb7b10b1e995eb53dec Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 11 Feb 2020 15:55:18 -0500
Subject: [PATCH 33/84] start

---
 tfjs-core/src/backends/cpu/softmax.ts | 56 +++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 tfjs-core/src/backends/cpu/softmax.ts

diff --git a/tfjs-core/src/backends/cpu/softmax.ts b/tfjs-core/src/backends/cpu/softmax.ts
new file mode 100644
index 00000000000..df77c5ffab3
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/softmax.ts
@@ -0,0 +1,56 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {max} from './kernel_names';
+import {ENGINE} from '../../engine';
+import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
+import {parseAxisParam} from '../../util';
+
+import {MathBackendCPU} from './backend_cpu';
+import {assertNotComplex} from './cpu_util';
+
+// import {max} from '../../ops/reduction_ops';
+
+interface SoftmaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface SoftmaxAttrs extends NamedAttrMap {
+  dim: number;
+}
+
+/**
+ * max_impl exports max kernel func
+ * max.ts uses register kernel, it imports max_impl
+ * softmax imports max_impl
+ */
+
+registerKernel({
+  kernelName: 'Softmax',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {x} = inputs as SoftmaxInputs;
+    const {dim} = attrs as SoftmaxAttrs;
+
+    const cpuBackend = backend as MathBackendCPU;
+    assertNotComplex(x, 'Softmax');
+
+    const axes = parseAxisParam([dim], x.shape);
+    // const maxLogit = max(x, axes);
+    const maxLogit = ENGINE.runKernel();
+  }
+});

From 2c9572bb744879d6b678bf565b68b106b980070b Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 11 Feb 2020 15:55:37 -0500
Subject: [PATCH 34/84] fix

---
 tfjs-backend-wasm/WORKSPACE    | 27 +++++++++++++--------------
 tfjs-backend-wasm/src/cc/BUILD |  2 +-
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/tfjs-backend-wasm/WORKSPACE b/tfjs-backend-wasm/WORKSPACE
index 88f02bc54f4..e76d39809be 100644
--- a/tfjs-backend-wasm/WORKSPACE
+++ b/tfjs-backend-wasm/WORKSPACE
@@ -16,72 +16,71 @@ git_repository(
 # The libraries below are transitive dependencies of XNNPACK that we need to
 # explicitly enumerate here. See https://docs.bazel.build/versions/master/external.html#transitive-dependencies
 
-
 # FP16 library, used for half-precision conversions
 http_archive(
     name = "FP16",
-    strip_prefix = "FP16-ba1d31f5eed2eb4a69e4dea3870a68c7c95f998f",
+    build_file = "@xnnpack//third_party:FP16.BUILD",
     sha256 = "9764297a339ad73b0717331a2c3e9c42a52105cd04cab62cb160e2b4598d2ea6",
+    strip_prefix = "FP16-ba1d31f5eed2eb4a69e4dea3870a68c7c95f998f",
     urls = [
         "https://github.com/Maratyszcza/FP16/archive/ba1d31f5eed2eb4a69e4dea3870a68c7c95f998f.tar.gz",
     ],
-    build_file = "@xnnpack//third_party:FP16.BUILD",
 )
 
 # FXdiv library, used for repeated integer division by the same factor
 http_archive(
     name = "FXdiv",
-    strip_prefix = "FXdiv-f8c5354679ec2597792bc70a9e06eff50c508b9a",
+    build_file = "@xnnpack//third_party:FXdiv.BUILD",
     sha256 = "7d3215bea832fe77091ec5666200b91156df6724da1e348205078346325fc45e",
+    strip_prefix = "FXdiv-f8c5354679ec2597792bc70a9e06eff50c508b9a",
     urls = [
         "https://github.com/Maratyszcza/FXdiv/archive/f8c5354679ec2597792bc70a9e06eff50c508b9a.tar.gz",
     ],
-    build_file = "@xnnpack//third_party:FXdiv.BUILD",
 )
 
 # pthreadpool library, used for parallelization
 http_archive(
     name = "pthreadpool",
+    build_file = "@xnnpack//third_party:pthreadpool.BUILD",
     sha256 = "c2328fdf9e48ac9b928953bcbc442eb14402d393e4cfae0541581a3d39efca9d",
     strip_prefix = "pthreadpool-0e275fe56094626349c55a524ea8b71a85daa64b",
     urls = [
-      "https://github.com/Maratyszcza/pthreadpool/archive/0e275fe56094626349c55a524ea8b71a85daa64b.tar.gz",
+        "https://github.com/Maratyszcza/pthreadpool/archive/0e275fe56094626349c55a524ea8b71a85daa64b.tar.gz",
     ],
-    build_file = "@xnnpack//third_party:pthreadpool.BUILD",
 )
 
 # clog library, used for logging
 http_archive(
     name = "clog",
-    strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
+    build_file = "@xnnpack//third_party:clog.BUILD",
     sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
+    strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
     urls = [
         "https://github.com/pytorch/cpuinfo/archive/d5e37adf1406cf899d7d9ec1d317c47506ccb970.tar.gz",
     ],
-    build_file = "@xnnpack//third_party:clog.BUILD",
 )
 
 # cpuinfo library, used for detecting processor characteristics
 http_archive(
     name = "cpuinfo",
-    strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
+    build_file = "@xnnpack//third_party:cpuinfo.BUILD",
+    patches = ["@xnnpack//third_party:cpuinfo.patch"],
     sha256 = "3f2dc1970f397a0e59db72f9fca6ff144b216895c1d606f6c94a507c1e53a025",
+    strip_prefix = "cpuinfo-d5e37adf1406cf899d7d9ec1d317c47506ccb970",
     urls = [
         "https://github.com/pytorch/cpuinfo/archive/d5e37adf1406cf899d7d9ec1d317c47506ccb970.tar.gz",
     ],
-    build_file = "@xnnpack//third_party:cpuinfo.BUILD",
-    patches = ["@xnnpack//third_party:cpuinfo.patch"],
 )
 
 # psimd library, used for fallback 128-bit SIMD micro-kernels
 http_archive(
     name = "psimd",
-    strip_prefix = "psimd-10b4ffc6ea9e2e11668f86969586f88bc82aaefa",
+    build_file = "@xnnpack//third_party:psimd.BUILD",
     sha256 = "1fefd66702cb2eb3462b962f33d4fb23d59a55d5889ee6372469d286c4512df4",
+    strip_prefix = "psimd-10b4ffc6ea9e2e11668f86969586f88bc82aaefa",
     urls = [
         "https://github.com/Maratyszcza/psimd/archive/10b4ffc6ea9e2e11668f86969586f88bc82aaefa.tar.gz",
     ],
-    build_file = "@xnnpack//third_party:psimd.BUILD",
 )
 
 git_repository(
diff --git a/tfjs-backend-wasm/src/cc/BUILD b/tfjs-backend-wasm/src/cc/BUILD
index e8e74ce1e52..21117098ccd 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -178,9 +178,9 @@ tfjs_cc_library(
         ":Minimum",
         ":Mul",
         ":Neg",
-        ":NotEqual",
         ":NonMaxSuppressionV3",
         ":NonMaxSuppressionV5",
+        ":NotEqual",
         ":PadV2",
         ":Prelu",
         ":Relu",

From 8d9f78925b282a0d937ad5bafb779b1a623dd191 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 06:46:05 -0500
Subject: [PATCH 35/84] edit

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  4 +-
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 66 +++++++++---------
 tfjs-core/src/backends/cpu/max.ts             | 63 +++++++++++++++++
 tfjs-core/src/backends/cpu/max_impl.ts        | 45 ++++++++++++
 tfjs-core/src/backends/cpu/softmax.ts         | 56 ---------------
 tfjs-core/src/backends/cpu/softmax_renamed.ts | 69 +++++++++++++++++++
 tfjs-core/src/ops/softmax_test.ts             |  7 +-
 7 files changed, 218 insertions(+), 92 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/max.ts
 create mode 100644 tfjs-core/src/backends/cpu/max_impl.ts
 delete mode 100644 tfjs-core/src/backends/cpu/softmax.ts
 create mode 100644 tfjs-core/src/backends/cpu/softmax_renamed.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index 75920b35b88..9eb32b3dfbc 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -18,5 +18,7 @@
 // We explicitly import the modular kernels so they get registered in the
 // global registry when we compile the library. A modular build would replace
 // the contents of this file and import only the kernels that are needed.
-import './square';
+import './max';
 import './non_max_suppression_v5';
+import './square';
+// import './softmax_renamed';
diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index 480770d17f4..d7d34dffd48 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,16 +377,16 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
-  softmax<T extends Tensor>(logits: T, dim: number): T {
-    const axes = util.parseAxisParam([dim], logits.shape);
-    const maxLogit = this.max(logits, axes);
-    const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
-    const a = this.subtract(logits, maxLogit.reshape(expandedShape));
-    const b = this.exp(a);
-    const sumExp = this.sum(b, axes).reshape(expandedShape);
+  // softmax<T extends Tensor>(logits: T, dim: number): T {
+  //   const axes = util.parseAxisParam([dim], logits.shape);
+  //   const maxLogit = this.max(logits, axes);
+  //   const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape,
+  //   axes); const a = this.subtract(logits, maxLogit.reshape(expandedShape));
+  //   const b = this.exp(a);
+  //   const sumExp = this.sum(b, axes).reshape(expandedShape);
 
-    return b.div(sumExp);
-  }
+  //   return b.div(sumExp);
+  // }
 
   subtract(a: Tensor, b: Tensor): Tensor {
     if (a.dtype === 'complex64' || b.dtype === 'complex64') {
@@ -826,30 +826,30 @@ export class MathBackendCPU extends KernelBackend {
     });
   }
 
-  max(x: Tensor, axes: number[]): Tensor {
-    assertNotComplex(x, 'max');
-
-    axis_util.assertAxesAreInnerMostDims('max', axes, x.rank);
-    const [outShape, reduceShape] =
-        axis_util.computeOutAndReduceShapes(x.shape, axes);
-    const result = ops.zeros(outShape, x.dtype);
-    const reduceSize = util.sizeFromShape(reduceShape);
-    const vals = this.readSync(result.dataId) as TypedArray;
-
-    const aVals = this.readSync(x.dataId) as TypedArray;
-    for (let i = 0; i < vals.length; ++i) {
-      const offset = i * reduceSize;
-      let max = aVals[offset];
-      for (let j = 0; j < reduceSize; ++j) {
-        const value = aVals[offset + j];
-        if (value > max) {
-          max = value;
-        }
-      }
-      vals[i] = max;
-    }
-    return result;
-  }
+  // max(x: Tensor, axes: number[]): Tensor {
+  //   assertNotComplex(x, 'max');
+
+  //   axis_util.assertAxesAreInnerMostDims('max', axes, x.rank);
+  //   const [outShape, reduceShape] =
+  //       axis_util.computeOutAndReduceShapes(x.shape, axes);
+  //   const result = ops.zeros(outShape, x.dtype);
+  //   const reduceSize = util.sizeFromShape(reduceShape);
+  //   const vals = this.readSync(result.dataId) as TypedArray;
+
+  //   const aVals = this.readSync(x.dataId) as TypedArray;
+  //   for (let i = 0; i < vals.length; ++i) {
+  //     const offset = i * reduceSize;
+  //     let max = aVals[offset];
+  //     for (let j = 0; j < reduceSize; ++j) {
+  //       const value = aVals[offset + j];
+  //       if (value > max) {
+  //         max = value;
+  //       }
+  //     }
+  //     vals[i] = max;
+  //   }
+  //   return result;
+  // }
 
   maximum(a: Tensor, b: Tensor): Tensor {
     assertNotComplex([a, b], 'maximum');
diff --git a/tfjs-core/src/backends/cpu/max.ts b/tfjs-core/src/backends/cpu/max.ts
new file mode 100644
index 00000000000..ac47b33f7ef
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/max.ts
@@ -0,0 +1,63 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {KernelFunc, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
+import * as axis_util from '../../ops/axis_util';
+import {sizeFromShape} from '../../util';
+
+import {MathBackendCPU} from './backend_cpu';
+import {assertNotComplex} from './cpu_util';
+
+interface MaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface MaxAttrs extends NamedAttrMap {
+  axes: number[];
+}
+
+export const max: KernelFunc = ({inputs, attrs, backend}) => {
+  const {x} = inputs as MaxInputs;
+  const {axes} = attrs as MaxAttrs;
+  const cpuBackend = backend as MathBackendCPU;
+  assertNotComplex(x, 'max');
+
+  axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+  const [outShape, reduceShape] =
+      axis_util.computeOutAndReduceShapes(x.shape, axes);
+
+  const aVals = cpuBackend.data.get(x.dataId).values as Float32Array;
+  const vals = new Float32Array(sizeFromShape(outShape));
+  const reduceSize = sizeFromShape(reduceShape);
+
+  for (let i = 0; i < vals.length; ++i) {
+    const offset = i * reduceSize;
+    let max = aVals[offset];
+    for (let j = 0; j < reduceSize; ++j) {
+      const value = aVals[offset + j];
+      if (value > max) {
+        max = value;
+      }
+    }
+    vals[i] = max;
+  }
+
+  const dataId = cpuBackend.write(vals, outShape, x.dtype);
+  return {dataId, shape: x.shape, dtype: x.dtype};
+};
+
+registerKernel({kernelName: 'Max', backendName: 'cpu', kernelFunc: max});
diff --git a/tfjs-core/src/backends/cpu/max_impl.ts b/tfjs-core/src/backends/cpu/max_impl.ts
new file mode 100644
index 00000000000..f1fb688c0a3
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/max_impl.ts
@@ -0,0 +1,45 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {TensorInfo} from '../../kernel_registry';
+import * as axis_util from '../../ops/axis_util';
+import {sizeFromShape} from '../../util';
+
+export const max = (x: TensorInfo, axes: number[]) => {
+  axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+  const [outShape, reduceShape] =
+      axis_util.computeOutAndReduceShapes(x.shape, axes);
+
+  const aVals = cpuBackend.data.get(x.dataId).values as Float32Array;
+  const vals = new Float32Array(sizeFromShape(outShape));
+  const reduceSize = sizeFromShape(reduceShape);
+
+  for (let i = 0; i < vals.length; ++i) {
+    const offset = i * reduceSize;
+    let max = aVals[offset];
+    for (let j = 0; j < reduceSize; ++j) {
+      const value = aVals[offset + j];
+      if (value > max) {
+        max = value;
+      }
+    }
+    vals[i] = max;
+  }
+
+  const dataId = cpuBackend.write(vals, outShape, x.dtype);
+  return {dataId, shape: x.shape, dtype: x.dtype};
+};
diff --git a/tfjs-core/src/backends/cpu/softmax.ts b/tfjs-core/src/backends/cpu/softmax.ts
deleted file mode 100644
index df77c5ffab3..00000000000
--- a/tfjs-core/src/backends/cpu/softmax.ts
+++ /dev/null
@@ -1,56 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {max} from './kernel_names';
-import {ENGINE} from '../../engine';
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
-import {parseAxisParam} from '../../util';
-
-import {MathBackendCPU} from './backend_cpu';
-import {assertNotComplex} from './cpu_util';
-
-// import {max} from '../../ops/reduction_ops';
-
-interface SoftmaxInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface SoftmaxAttrs extends NamedAttrMap {
-  dim: number;
-}
-
-/**
- * max_impl exports max kernel func
- * max.ts uses register kernel, it imports max_impl
- * softmax imports max_impl
- */
-
-registerKernel({
-  kernelName: 'Softmax',
-  backendName: 'cpu',
-  kernelFunc: ({inputs, attrs, backend}) => {
-    const {x} = inputs as SoftmaxInputs;
-    const {dim} = attrs as SoftmaxAttrs;
-
-    const cpuBackend = backend as MathBackendCPU;
-    assertNotComplex(x, 'Softmax');
-
-    const axes = parseAxisParam([dim], x.shape);
-    // const maxLogit = max(x, axes);
-    const maxLogit = ENGINE.runKernel();
-  }
-});
diff --git a/tfjs-core/src/backends/cpu/softmax_renamed.ts b/tfjs-core/src/backends/cpu/softmax_renamed.ts
new file mode 100644
index 00000000000..b645497d06f
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/softmax_renamed.ts
@@ -0,0 +1,69 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {max} from './kernel_names';
+// import {ENGINE} from '../../engine';
+// import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from
+// '../../kernel_registry';
+import {KernelFunc, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
+import {parseAxisParam} from '../../util';
+
+import {MathBackendCPU} from './backend_cpu';
+import {assertNotComplex} from './cpu_util';
+import {max} from './max';
+
+interface SoftmaxInputs extends NamedTensorInfoMap {
+  logits: TensorInfo;
+}
+
+interface SoftmaxAttrs extends NamedAttrMap {
+  dim: number;
+}
+
+/**
+ * max_impl exports max kernel func
+ * max.ts uses register kernel, it imports max_impl
+ * softmax imports max_impl
+ */
+
+export const softmax: KernelFunc = ({inputs, attrs, backend}) => {
+  const {logits} = inputs as SoftmaxInputs;
+  const {dim} = attrs as SoftmaxAttrs;
+
+  const cpuBackend = backend as MathBackendCPU;
+  assertNotComplex(logits, 'Softmax');
+  console.log('SOFTMAXJasf sdf');
+  console.log(logits);
+
+  const axes = parseAxisParam([dim], logits.shape);
+  const maxLogit = max(logits, axes);
+  console.log(maxLogit);
+  // const maxLogit = ENGINE.runKernel();
+
+  const values = cpuBackend.data.get(logits.dataId).values as Float32Array;
+  const newValues = new Float32Array(values.length);
+  for (let i = 0; i < values.length; ++i) {
+    const value = values[i];
+    newValues[i] = value;
+  }
+
+  const dataId = cpuBackend.write(newValues, logits.shape, logits.dtype);
+  return {dataId, shape: logits.shape, dtype: 'float32'};
+};
+
+registerKernel(
+    {kernelName: 'Softmax', backendName: 'cpu', kernelFunc: softmax});
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index e18397bb0a0..9915fca14a7 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,18 +15,21 @@
  * =============================================================================
  */
 
+import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', ALL_ENVS, () => {
-  it('regular test', async () => {
+describeWithFlags('softmax', CPU_ENVS, () => {
+  fit('regular test', async () => {
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
     expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
     expectArraysClose(await y.sum().data(), 1);
   });
+});
 
+describeWithFlags('softmax', ALL_ENVS, () => {
   it('overflow', async () => {
     const y = tf.softmax(tf.tensor1d([100, 100]));
 

From 29a7cdc03b433c6e5b2ec299b5cd2eeff0be03ca Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 10:38:53 -0500
Subject: [PATCH 36/84] remove

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  2 -
 tfjs-core/src/backends/cpu/max.ts             | 63 -----------------
 tfjs-core/src/backends/cpu/max_impl.ts        | 45 ------------
 tfjs-core/src/backends/cpu/softmax_renamed.ts | 69 -------------------
 tfjs-core/src/ops/softmax_test.ts             |  2 +-
 5 files changed, 1 insertion(+), 180 deletions(-)
 delete mode 100644 tfjs-core/src/backends/cpu/max.ts
 delete mode 100644 tfjs-core/src/backends/cpu/max_impl.ts
 delete mode 100644 tfjs-core/src/backends/cpu/softmax_renamed.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index 9eb32b3dfbc..b9191e1b4c2 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -18,7 +18,5 @@
 // We explicitly import the modular kernels so they get registered in the
 // global registry when we compile the library. A modular build would replace
 // the contents of this file and import only the kernels that are needed.
-import './max';
 import './non_max_suppression_v5';
 import './square';
-// import './softmax_renamed';
diff --git a/tfjs-core/src/backends/cpu/max.ts b/tfjs-core/src/backends/cpu/max.ts
deleted file mode 100644
index ac47b33f7ef..00000000000
--- a/tfjs-core/src/backends/cpu/max.ts
+++ /dev/null
@@ -1,63 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {KernelFunc, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
-import * as axis_util from '../../ops/axis_util';
-import {sizeFromShape} from '../../util';
-
-import {MathBackendCPU} from './backend_cpu';
-import {assertNotComplex} from './cpu_util';
-
-interface MaxInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface MaxAttrs extends NamedAttrMap {
-  axes: number[];
-}
-
-export const max: KernelFunc = ({inputs, attrs, backend}) => {
-  const {x} = inputs as MaxInputs;
-  const {axes} = attrs as MaxAttrs;
-  const cpuBackend = backend as MathBackendCPU;
-  assertNotComplex(x, 'max');
-
-  axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
-  const [outShape, reduceShape] =
-      axis_util.computeOutAndReduceShapes(x.shape, axes);
-
-  const aVals = cpuBackend.data.get(x.dataId).values as Float32Array;
-  const vals = new Float32Array(sizeFromShape(outShape));
-  const reduceSize = sizeFromShape(reduceShape);
-
-  for (let i = 0; i < vals.length; ++i) {
-    const offset = i * reduceSize;
-    let max = aVals[offset];
-    for (let j = 0; j < reduceSize; ++j) {
-      const value = aVals[offset + j];
-      if (value > max) {
-        max = value;
-      }
-    }
-    vals[i] = max;
-  }
-
-  const dataId = cpuBackend.write(vals, outShape, x.dtype);
-  return {dataId, shape: x.shape, dtype: x.dtype};
-};
-
-registerKernel({kernelName: 'Max', backendName: 'cpu', kernelFunc: max});
diff --git a/tfjs-core/src/backends/cpu/max_impl.ts b/tfjs-core/src/backends/cpu/max_impl.ts
deleted file mode 100644
index f1fb688c0a3..00000000000
--- a/tfjs-core/src/backends/cpu/max_impl.ts
+++ /dev/null
@@ -1,45 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {TensorInfo} from '../../kernel_registry';
-import * as axis_util from '../../ops/axis_util';
-import {sizeFromShape} from '../../util';
-
-export const max = (x: TensorInfo, axes: number[]) => {
-  axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
-  const [outShape, reduceShape] =
-      axis_util.computeOutAndReduceShapes(x.shape, axes);
-
-  const aVals = cpuBackend.data.get(x.dataId).values as Float32Array;
-  const vals = new Float32Array(sizeFromShape(outShape));
-  const reduceSize = sizeFromShape(reduceShape);
-
-  for (let i = 0; i < vals.length; ++i) {
-    const offset = i * reduceSize;
-    let max = aVals[offset];
-    for (let j = 0; j < reduceSize; ++j) {
-      const value = aVals[offset + j];
-      if (value > max) {
-        max = value;
-      }
-    }
-    vals[i] = max;
-  }
-
-  const dataId = cpuBackend.write(vals, outShape, x.dtype);
-  return {dataId, shape: x.shape, dtype: x.dtype};
-};
diff --git a/tfjs-core/src/backends/cpu/softmax_renamed.ts b/tfjs-core/src/backends/cpu/softmax_renamed.ts
deleted file mode 100644
index b645497d06f..00000000000
--- a/tfjs-core/src/backends/cpu/softmax_renamed.ts
+++ /dev/null
@@ -1,69 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {max} from './kernel_names';
-// import {ENGINE} from '../../engine';
-// import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from
-// '../../kernel_registry';
-import {KernelFunc, NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
-import {parseAxisParam} from '../../util';
-
-import {MathBackendCPU} from './backend_cpu';
-import {assertNotComplex} from './cpu_util';
-import {max} from './max';
-
-interface SoftmaxInputs extends NamedTensorInfoMap {
-  logits: TensorInfo;
-}
-
-interface SoftmaxAttrs extends NamedAttrMap {
-  dim: number;
-}
-
-/**
- * max_impl exports max kernel func
- * max.ts uses register kernel, it imports max_impl
- * softmax imports max_impl
- */
-
-export const softmax: KernelFunc = ({inputs, attrs, backend}) => {
-  const {logits} = inputs as SoftmaxInputs;
-  const {dim} = attrs as SoftmaxAttrs;
-
-  const cpuBackend = backend as MathBackendCPU;
-  assertNotComplex(logits, 'Softmax');
-  console.log('SOFTMAXJasf sdf');
-  console.log(logits);
-
-  const axes = parseAxisParam([dim], logits.shape);
-  const maxLogit = max(logits, axes);
-  console.log(maxLogit);
-  // const maxLogit = ENGINE.runKernel();
-
-  const values = cpuBackend.data.get(logits.dataId).values as Float32Array;
-  const newValues = new Float32Array(values.length);
-  for (let i = 0; i < values.length; ++i) {
-    const value = values[i];
-    newValues[i] = value;
-  }
-
-  const dataId = cpuBackend.write(newValues, logits.shape, logits.dtype);
-  return {dataId, shape: logits.shape, dtype: 'float32'};
-};
-
-registerKernel(
-    {kernelName: 'Softmax', backendName: 'cpu', kernelFunc: softmax});
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 9915fca14a7..6bbb9df3068 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -21,7 +21,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', CPU_ENVS, () => {
-  fit('regular test', async () => {
+  it('regular test', async () => {
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
     expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);

From 2ae0a59fca0967181cf271281384fadaf6757e90 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 10:43:14 -0500
Subject: [PATCH 37/84] revive

---
 tfjs-core/src/backends/backend.ts             |  2 +-
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 66 +++++++++----------
 tfjs-core/src/backends/webgl/backend_webgl.ts |  4 +-
 3 files changed, 36 insertions(+), 36 deletions(-)

diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts
index d5e9cd02b4b..abfbe4fc52f 100644
--- a/tfjs-core/src/backends/backend.ts
+++ b/tfjs-core/src/backends/backend.ts
@@ -305,7 +305,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
   expm1<T extends Tensor>(x: T): T {
     return notYetImplemented('expm1');
   }
-  softmax<T extends Tensor>(x: T, dim: number): T {
+  softmax(x: Tensor, dim: number): Tensor {
     return notYetImplemented('softmax');
   }
   log<T extends Tensor>(x: T): T {
diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index d7d34dffd48..1982d41bcce 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,16 +377,16 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
-  // softmax<T extends Tensor>(logits: T, dim: number): T {
-  //   const axes = util.parseAxisParam([dim], logits.shape);
-  //   const maxLogit = this.max(logits, axes);
-  //   const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape,
-  //   axes); const a = this.subtract(logits, maxLogit.reshape(expandedShape));
-  //   const b = this.exp(a);
-  //   const sumExp = this.sum(b, axes).reshape(expandedShape);
+  softmax(logits: Tensor, dim: number): Tensor {
+    const axes = util.parseAxisParam([dim], logits.shape);
+    const maxLogit = this.max(logits, axes);
+    const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
+    const a = this.subtract(logits, maxLogit.reshape(expandedShape));
+    const b = this.exp(a);
+    const sumExp = this.sum(b, axes).reshape(expandedShape);
 
-  //   return b.div(sumExp);
-  // }
+    return this.realDivide(b, sumExp);
+  }
 
   subtract(a: Tensor, b: Tensor): Tensor {
     if (a.dtype === 'complex64' || b.dtype === 'complex64') {
@@ -826,30 +826,30 @@ export class MathBackendCPU extends KernelBackend {
     });
   }
 
-  // max(x: Tensor, axes: number[]): Tensor {
-  //   assertNotComplex(x, 'max');
-
-  //   axis_util.assertAxesAreInnerMostDims('max', axes, x.rank);
-  //   const [outShape, reduceShape] =
-  //       axis_util.computeOutAndReduceShapes(x.shape, axes);
-  //   const result = ops.zeros(outShape, x.dtype);
-  //   const reduceSize = util.sizeFromShape(reduceShape);
-  //   const vals = this.readSync(result.dataId) as TypedArray;
-
-  //   const aVals = this.readSync(x.dataId) as TypedArray;
-  //   for (let i = 0; i < vals.length; ++i) {
-  //     const offset = i * reduceSize;
-  //     let max = aVals[offset];
-  //     for (let j = 0; j < reduceSize; ++j) {
-  //       const value = aVals[offset + j];
-  //       if (value > max) {
-  //         max = value;
-  //       }
-  //     }
-  //     vals[i] = max;
-  //   }
-  //   return result;
-  // }
+  max(x: Tensor, axes: number[]): Tensor {
+    assertNotComplex(x, 'max');
+
+    axis_util.assertAxesAreInnerMostDims('max', axes, x.rank);
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(x.shape, axes);
+    const result = ops.zeros(outShape, x.dtype);
+    const reduceSize = util.sizeFromShape(reduceShape);
+    const vals = this.readSync(result.dataId) as TypedArray;
+
+    const aVals = this.readSync(x.dataId) as TypedArray;
+    for (let i = 0; i < vals.length; ++i) {
+      const offset = i * reduceSize;
+      let max = aVals[offset];
+      for (let j = 0; j < reduceSize; ++j) {
+        const value = aVals[offset + j];
+        if (value > max) {
+          max = value;
+        }
+      }
+      vals[i] = max;
+    }
+    return result;
+  }
 
   maximum(a: Tensor, b: Tensor): Tensor {
     assertNotComplex([a, b], 'maximum');
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 151324aea99..6f9b832e09e 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1597,7 +1597,7 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  softmax<T extends Tensor>(logits: T, dim: number): T {
+  softmax(logits: Tensor, dim: number): Tensor {
     const axes = util.parseAxisParam([dim], logits.shape);
     const maxLogit = this.max(logits, axes);
     const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
@@ -1605,7 +1605,7 @@ export class MathBackendWebGL extends KernelBackend {
     const b = this.exp(a);
     const sumExp = this.sum(b, axes).reshape(expandedShape);
 
-    return b.div(sumExp);
+    return this.realDivide(b, sumExp);
   }
 
   log<T extends Tensor>(x: T): T {

From 4442f314494f1b345fa1e8c0e8b2b13dd18dae21 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 10:47:28 -0500
Subject: [PATCH 38/84] revive

---
 tfjs-core/src/ops/softmax_test.ts | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 6bbb9df3068..e18397bb0a0 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,21 +15,18 @@
  * =============================================================================
  */
 
-import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', CPU_ENVS, () => {
+describeWithFlags('softmax', ALL_ENVS, () => {
   it('regular test', async () => {
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
     expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
     expectArraysClose(await y.sum().data(), 1);
   });
-});
 
-describeWithFlags('softmax', ALL_ENVS, () => {
   it('overflow', async () => {
     const y = tf.softmax(tf.tensor1d([100, 100]));
 

From c4385c44796d9301143d8d5d30d6267b69c71044 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 11:02:41 -0500
Subject: [PATCH 39/84] build

---
 tfjs-core/src/backends/backend.ts             | 2 +-
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 4 ++--
 tfjs-core/src/backends/webgl/backend_webgl.ts | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tfjs-core/src/backends/backend.ts b/tfjs-core/src/backends/backend.ts
index abfbe4fc52f..d5e9cd02b4b 100644
--- a/tfjs-core/src/backends/backend.ts
+++ b/tfjs-core/src/backends/backend.ts
@@ -305,7 +305,7 @@ export class KernelBackend implements TensorStorage, Backend, BackendTimer {
   expm1<T extends Tensor>(x: T): T {
     return notYetImplemented('expm1');
   }
-  softmax(x: Tensor, dim: number): Tensor {
+  softmax<T extends Tensor>(x: T, dim: number): T {
     return notYetImplemented('softmax');
   }
   log<T extends Tensor>(x: T): T {
diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index 1982d41bcce..f45855309a8 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,7 +377,7 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
-  softmax(logits: Tensor, dim: number): Tensor {
+  softmax<T extends Tensor>(logits: T, dim: number): T {
     const axes = util.parseAxisParam([dim], logits.shape);
     const maxLogit = this.max(logits, axes);
     const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
@@ -385,7 +385,7 @@ export class MathBackendCPU extends KernelBackend {
     const b = this.exp(a);
     const sumExp = this.sum(b, axes).reshape(expandedShape);
 
-    return this.realDivide(b, sumExp);
+    return this.realDivide(b, sumExp) as T;
   }
 
   subtract(a: Tensor, b: Tensor): Tensor {
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 6f9b832e09e..a1dfe5c0086 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1597,7 +1597,7 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  softmax(logits: Tensor, dim: number): Tensor {
+  softmax<T extends Tensor>(logits: T, dim: number): T {
     const axes = util.parseAxisParam([dim], logits.shape);
     const maxLogit = this.max(logits, axes);
     const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
@@ -1605,7 +1605,7 @@ export class MathBackendWebGL extends KernelBackend {
     const b = this.exp(a);
     const sumExp = this.sum(b, axes).reshape(expandedShape);
 
-    return this.realDivide(b, sumExp);
+    return this.realDivide(b, sumExp) as T;
   }
 
   log<T extends Tensor>(x: T): T {

From 82921ca54662ffb15e01f4aae4fc7d0a3d0bd684 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 11:11:34 -0500
Subject: [PATCH 40/84] revive test

---
 tfjs-core/src/ops/softmax_test.ts | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index e18397bb0a0..9915fca14a7 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,18 +15,21 @@
  * =============================================================================
  */
 
+import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', ALL_ENVS, () => {
-  it('regular test', async () => {
+describeWithFlags('softmax', CPU_ENVS, () => {
+  fit('regular test', async () => {
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
     expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
     expectArraysClose(await y.sum().data(), 1);
   });
+});
 
+describeWithFlags('softmax', ALL_ENVS, () => {
   it('overflow', async () => {
     const y = tf.softmax(tf.tensor1d([100, 100]));
 

From cde91358371bf49b3ccd168a67f3dfef6f72d4ba Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 13:02:27 -0500
Subject: [PATCH 41/84] move kernels

---
 tfjs-core/src/backends/cpu/backend_cpu.ts          | 1 +
 tfjs-core/src/backends/cpu/kernels/softmax.ts      | 0
 tfjs-core/src/backends/cpu/{ => kernels}/square.ts | 6 +++---
 3 files changed, 4 insertions(+), 3 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/softmax.ts
 rename tfjs-core/src/backends/cpu/{ => kernels}/square.ts (92%)

diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index f45855309a8..a94a5b6673b 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -378,6 +378,7 @@ export class MathBackendCPU extends KernelBackend {
   }
 
   softmax<T extends Tensor>(logits: T, dim: number): T {
+    console.log('softmaxxxxx');
     const axes = util.parseAxisParam([dim], logits.shape);
     const maxLogit = this.max(logits, axes);
     const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tfjs-core/src/backends/cpu/square.ts b/tfjs-core/src/backends/cpu/kernels/square.ts
similarity index 92%
rename from tfjs-core/src/backends/cpu/square.ts
rename to tfjs-core/src/backends/cpu/kernels/square.ts
index b07b15f76a0..9e613aebaf0 100644
--- a/tfjs-core/src/backends/cpu/square.ts
+++ b/tfjs-core/src/backends/cpu/kernels/square.ts
@@ -15,10 +15,10 @@
  * =============================================================================
  */
 
-import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../kernel_registry';
+import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 
-import {MathBackendCPU} from './backend_cpu';
-import {assertNotComplex} from './cpu_util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
 
 interface SquareInputs extends NamedTensorInfoMap {
   x: TensorInfo;

From 657f5c53a051c93723a6e0cdf68a14e6e8dc8ae7 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 15:38:28 -0500
Subject: [PATCH 42/84] start

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  4 +-
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 22 +++----
 tfjs-core/src/backends/cpu/kernels/max.ts     | 55 ++++++++++++++++++
 .../src/backends/cpu/kernels/max_impl.ts      | 39 +++++++++++++
 tfjs-core/src/backends/cpu/kernels/softmax.ts | 58 +++++++++++++++++++
 5 files changed, 166 insertions(+), 12 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/max.ts
 create mode 100644 tfjs-core/src/backends/cpu/kernels/max_impl.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index b9191e1b4c2..88470600bcc 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -19,4 +19,6 @@
 // global registry when we compile the library. A modular build would replace
 // the contents of this file and import only the kernels that are needed.
 import './non_max_suppression_v5';
-import './square';
+import './kernels/square';
+import './kernels/softmax';
+import './kernels/max';
diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index a94a5b6673b..af5cda78a23 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,17 +377,17 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
-  softmax<T extends Tensor>(logits: T, dim: number): T {
-    console.log('softmaxxxxx');
-    const axes = util.parseAxisParam([dim], logits.shape);
-    const maxLogit = this.max(logits, axes);
-    const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
-    const a = this.subtract(logits, maxLogit.reshape(expandedShape));
-    const b = this.exp(a);
-    const sumExp = this.sum(b, axes).reshape(expandedShape);
-
-    return this.realDivide(b, sumExp) as T;
-  }
+  // softmax<T extends Tensor>(logits: T, dim: number): T {
+  //   console.log('softmaxxxxx');
+  //   const axes = util.parseAxisParam([dim], logits.shape);
+  //   const maxLogit = this.max(logits, axes);
+  //   const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape,
+  //   axes); const a = this.subtract(logits, maxLogit.reshape(expandedShape));
+  //   const b = this.exp(a);
+  //   const sumExp = this.sum(b, axes).reshape(expandedShape);
+
+  //   return this.realDivide(b, sumExp) as T;
+  // }
 
   subtract(a: Tensor, b: Tensor): Tensor {
     if (a.dtype === 'complex64' || b.dtype === 'complex64') {
diff --git a/tfjs-core/src/backends/cpu/kernels/max.ts b/tfjs-core/src/backends/cpu/kernels/max.ts
new file mode 100644
index 00000000000..2a687b19284
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/max.ts
@@ -0,0 +1,55 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. 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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+import {max} from './max_impl';
+
+interface MaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface MaxAttrs extends NamedAttrMap {
+  axes: number[];
+}
+
+registerKernel({
+  kernelName: 'Max',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {x} = inputs as MaxInputs;
+    const {axes} = attrs as MaxAttrs;
+    const cpuBackend = backend as MathBackendCPU;
+
+    assertNotComplex(x, 'max');
+
+    axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(x.shape, axes);
+
+    const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
+    const outValues = new Float32Array(sizeFromShape(outShape));
+    const result = max(xVals, reduceShape, outValues);
+
+    const dataId = cpuBackend.write(result, outShape, x.dtype);
+    return {dataId, shape: outShape, dtype: x.dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/cpu/kernels/max_impl.ts b/tfjs-core/src/backends/cpu/kernels/max_impl.ts
new file mode 100644
index 00000000000..49eee5d4ce2
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/max_impl.ts
@@ -0,0 +1,39 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {TypedArray} from '../../../types';
+import {sizeFromShape} from '../../../util';
+
+export const max =
+    (x: TypedArray, reduceShape: number[], outValues: TypedArray):
+        TypedArray => {
+          const reduceSize = sizeFromShape(reduceShape);
+
+          for (let i = 0; i < x.length; ++i) {
+            const offset = i * reduceSize;
+            let max = x[offset];
+            for (let j = 0; j < reduceSize; ++j) {
+              const value = x[offset + j];
+              if (value > max) {
+                max = value;
+              }
+            }
+            outValues[i] = max;
+          }
+
+          return outValues;
+        };
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index e69de29bb2d..ff8eb5b85fd 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -0,0 +1,58 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. 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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {parseAxisParam, sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+import {max} from './max_impl';
+
+interface SoftmaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface SoftmaxAttrs extends NamedAttrMap {
+  dim: number;
+}
+
+registerKernel({
+  kernelName: 'Softmax',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    console.log('BLERG KERNEL FUNC');
+    const {logits} = inputs as SoftmaxInputs;
+    const {dim} = attrs as SoftmaxAttrs;
+    const cpuBackend = backend as MathBackendCPU;
+    assertNotComplex(logits, 'softmax');
+
+    const axes = parseAxisParam([dim], logits.shape);
+
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(logits.shape, axes);
+    const values = cpuBackend.data.get(logits.dataId).values as Float32Array;
+    const outValues = new Float32Array(sizeFromShape(outShape));
+    const maxLogit = max(values, reduceShape, outValues);
+
+    console.log('MAX LOGIT');
+    console.log(maxLogit);
+
+    const dataId = cpuBackend.write(maxLogit, outShape, logits.dtype);
+    return {dataId, shape: outShape, dtype: logits.dtype};
+  }
+});

From 3cda765735c9e4244648cfcd0f43da45e4e5aee6 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 15:45:35 -0500
Subject: [PATCH 43/84] add sub

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  1 +
 tfjs-core/src/backends/cpu/kernels/Sub.ts     | 51 +++++++++++++++++++
 tfjs-core/src/backends/cpu/kernels/softmax.ts | 11 ++--
 3 files changed, 59 insertions(+), 4 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/Sub.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index 88470600bcc..49a5ef26aa9 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -22,3 +22,4 @@ import './non_max_suppression_v5';
 import './kernels/square';
 import './kernels/softmax';
 import './kernels/max';
+import './kernels/Sub';
diff --git a/tfjs-core/src/backends/cpu/kernels/Sub.ts b/tfjs-core/src/backends/cpu/kernels/Sub.ts
new file mode 100644
index 00000000000..536b972a8d0
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/Sub.ts
@@ -0,0 +1,51 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+import {max} from './max_impl';
+
+interface SubInputs extends NamedTensorInfoMap {
+  a: TensorInfo;
+  b: TensorInfo;
+}
+
+registerKernel({
+  kernelName: 'Sub',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend}) => {
+    const {a, b} = inputs as SubInputs;
+    const cpuBackend = backend as MathBackendCPU;
+
+    assertNotComplex(x, 'max');
+
+    axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(x.shape, axes);
+
+    const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
+    const outValues = new Float32Array(sizeFromShape(outShape));
+    const result = max(xVals, reduceShape, outValues);
+
+    const dataId = cpuBackend.write(result, outShape, x.dtype);
+    return {dataId, shape: outShape, dtype: x.dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index ff8eb5b85fd..1a206716487 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -43,16 +43,19 @@ registerKernel({
 
     const axes = parseAxisParam([dim], logits.shape);
 
-    const [outShape, reduceShape] =
+    const [maxLogitOutShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
     const values = cpuBackend.data.get(logits.dataId).values as Float32Array;
-    const outValues = new Float32Array(sizeFromShape(outShape));
+    const outValues = new Float32Array(sizeFromShape(maxLogitOutShape));
     const maxLogit = max(values, reduceShape, outValues);
 
+    const expandedShape =
+        axis_util.expandShapeToKeepDim(maxLogitOutShape, axes);
+
     console.log('MAX LOGIT');
     console.log(maxLogit);
 
-    const dataId = cpuBackend.write(maxLogit, outShape, logits.dtype);
-    return {dataId, shape: outShape, dtype: logits.dtype};
+    const dataId = cpuBackend.write(maxLogit, maxLogitOutShape, logits.dtype);
+    return {dataId, shape: maxLogitOutShape, dtype: logits.dtype};
   }
 });

From 5991de392a5e94efbee3bed0a780ae0979598540 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 09:47:05 -0500
Subject: [PATCH 44/84] subimpl

---
 tfjs-core/src/backends/cpu/kernels/Sub.ts     | 26 +++++-----
 tfjs-core/src/backends/cpu/kernels/softmax.ts | 17 ++++--
 .../src/backends/cpu/kernels/sub_impl.ts      | 52 +++++++++++++++++++
 tfjs-core/src/util.ts                         | 32 ++++++++++++
 4 files changed, 110 insertions(+), 17 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/sub_impl.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/Sub.ts b/tfjs-core/src/backends/cpu/kernels/Sub.ts
index 536b972a8d0..fa3bb841235 100644
--- a/tfjs-core/src/backends/cpu/kernels/Sub.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Sub.ts
@@ -16,11 +16,12 @@
  */
 
 import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
-import * as axis_util from '../../../ops/axis_util';
+import * as broadcast_util from '../../../ops/broadcast_util';
+import {upcastType} from '../../../types';
 import {sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
-import {assertNotComplex} from '../cpu_util';
-import {max} from './max_impl';
+
+import {sub} from './sub_impl';
 
 interface SubInputs extends NamedTensorInfoMap {
   a: TensorInfo;
@@ -34,18 +35,17 @@ registerKernel({
     const {a, b} = inputs as SubInputs;
     const cpuBackend = backend as MathBackendCPU;
 
-    assertNotComplex(x, 'max');
-
-    axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+    const dtype = upcastType(a.dtype, b.dtype);
+    const outShape =
+        broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
+    const outValues = new Float32Array(sizeFromShape(outShape));
 
-    const [outShape, reduceShape] =
-        axis_util.computeOutAndReduceShapes(x.shape, axes);
+    const aVals = cpuBackend.data.get(a.dataId).values as Float32Array;
+    const bVals = cpuBackend.data.get(b.dataId).values as Float32Array;
 
-    const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
-    const outValues = new Float32Array(sizeFromShape(outShape));
-    const result = max(xVals, reduceShape, outValues);
+    const result = sub(aVals, a.shape, bVals, b.shape, outValues, outShape);
 
-    const dataId = cpuBackend.write(result, outShape, x.dtype);
-    return {dataId, shape: outShape, dtype: x.dtype};
+    const dataId = cpuBackend.write(result, outShape, dtype);
+    return {dataId, shape: outShape, dtype};
   }
 });
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index 1a206716487..8b0f5f05457 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -22,6 +22,7 @@ import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
 import {max} from './max_impl';
+import {sub} from './sub_impl';
 
 interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
@@ -35,7 +36,6 @@ registerKernel({
   kernelName: 'Softmax',
   backendName: 'cpu',
   kernelFunc: ({inputs, attrs, backend}) => {
-    console.log('BLERG KERNEL FUNC');
     const {logits} = inputs as SoftmaxInputs;
     const {dim} = attrs as SoftmaxAttrs;
     const cpuBackend = backend as MathBackendCPU;
@@ -45,9 +45,11 @@ registerKernel({
 
     const [maxLogitOutShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
-    const values = cpuBackend.data.get(logits.dataId).values as Float32Array;
-    const outValues = new Float32Array(sizeFromShape(maxLogitOutShape));
-    const maxLogit = max(values, reduceShape, outValues);
+    const logitsValues =
+        cpuBackend.data.get(logits.dataId).values as Float32Array;
+    const maxLogit =
+        max(logitsValues, reduceShape,
+            new Float32Array(sizeFromShape(maxLogitOutShape)));
 
     const expandedShape =
         axis_util.expandShapeToKeepDim(maxLogitOutShape, axes);
@@ -55,6 +57,13 @@ registerKernel({
     console.log('MAX LOGIT');
     console.log(maxLogit);
 
+    const a =
+        sub(logitsValues, logits.shape, maxLogit, expandedShape,
+            new Float32Array(sizeFromShape(logits.shape)), logits.shape);
+
+    console.log('subtract');
+    console.log(a);
+
     const dataId = cpuBackend.write(maxLogit, maxLogitOutShape, logits.dtype);
     return {dataId, shape: maxLogitOutShape, dtype: logits.dtype};
   }
diff --git a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
new file mode 100644
index 00000000000..a6a6b34335d
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 broadcast_util from '../../../ops/broadcast_util';
+import {TypedArray} from '../../../types';
+import * as util from '../../../util';
+
+export const sub =
+    (a: TypedArray, aShape: number[], b: TypedArray, bShape: number[],
+     outValues: TypedArray, outShape: number[]): TypedArray => {
+      const aBroadcastDims = broadcast_util.getBroadcastDims(aShape, outShape);
+      const bBroadcastDims = broadcast_util.getBroadcastDims(bShape, outShape);
+
+      const aStrides = util.computeStrides(aShape);
+      const bStrides = util.computeStrides(bShape);
+      const outStrides = util.computeStrides(outShape);
+
+      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
+        for (let i = 0; i < outValues.length; ++i) {
+          outValues[i] = a[i % a.length] - b[i % b.length];
+        }
+      } else {
+        for (let i = 0; i < outValues.length; ++i) {
+          const loc = util.indexToLoc(i, outStrides);
+
+          const aLoc = loc.slice(-aShape.length);
+          aBroadcastDims.forEach(d => aLoc[d] = 0);
+          const aIndex = util.locToIndex(aLoc, aStrides);
+
+          const bLoc = loc.slice(-bShape.length);
+          bBroadcastDims.forEach(d => bLoc[d] = 0);
+          const bIndex = util.locToIndex(bLoc, bStrides);
+
+          outValues[i] = a[aIndex] - b[bIndex];
+        }
+      }
+      return outValues;
+    };
diff --git a/tfjs-core/src/util.ts b/tfjs-core/src/util.ts
index 8edc2013427..7c213f527f9 100644
--- a/tfjs-core/src/util.ts
+++ b/tfjs-core/src/util.ts
@@ -548,6 +548,38 @@ export function computeStrides(shape: number[]): number[] {
   return strides;
 }
 
+export function locToIndex(locs: number[], strides: number[]): number {
+  const rank = locs.length;
+  if (rank === 0) {
+    return 0;
+  }
+  if (rank === 1) {
+    return locs[0];
+  }
+  let index = locs[rank - 1];
+  for (let i = 0; i < rank - 1; ++i) {
+    index += strides[i] * locs[i];
+  }
+  return index;
+}
+
+export function indexToLoc(index: number, strides: number[]): number[] {
+  const rank = strides.length + 1;
+  if (rank === 0) {
+    return [];
+  }
+  if (rank === 1) {
+    return [index];
+  }
+  const locs = new Array(rank);
+  for (let i = 0; i < locs.length - 1; ++i) {
+    locs[i] = Math.floor(index / strides[i]);
+    index -= locs[i] * strides[i];
+  }
+  locs[locs.length - 1] = index;
+  return locs;
+}
+
 export function toTypedArray(
     a: TensorLike, dtype: DataType, debugMode: boolean): TypedArray {
   if (dtype === 'string') {

From bdddd58cca9171f98f23c970d60bcf0c94367648 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 10:18:08 -0500
Subject: [PATCH 45/84] add exp

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  1 +
 tfjs-core/src/backends/cpu/kernels/Exp.ts     | 42 +++++++++++++++++++
 .../src/backends/cpu/kernels/exp_impl.ts      | 26 ++++++++++++
 tfjs-core/src/backends/cpu/kernels/softmax.ts |  5 +++
 4 files changed, 74 insertions(+)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/Exp.ts
 create mode 100644 tfjs-core/src/backends/cpu/kernels/exp_impl.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index 49a5ef26aa9..ef85b433af4 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -23,3 +23,4 @@ import './kernels/square';
 import './kernels/softmax';
 import './kernels/max';
 import './kernels/Sub';
+import './kernels/Exp';
diff --git a/tfjs-core/src/backends/cpu/kernels/Exp.ts b/tfjs-core/src/backends/cpu/kernels/Exp.ts
new file mode 100644
index 00000000000..e834aa6e1fd
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/Exp.ts
@@ -0,0 +1,42 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+
+import {exp} from './exp_impl';
+
+interface ExpInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+registerKernel({
+  kernelName: 'Exp',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend}) => {
+    const {x} = inputs as ExpInputs;
+    const cpuBackend = backend as MathBackendCPU;
+
+    const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
+
+    const result = exp(xVals, new Float32Array(sizeFromShape(x.shape)));
+
+    const dataId = cpuBackend.write(result, x.shape, x.dtype);
+    return {dataId, shape: x.shape, dtype: x.dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/cpu/kernels/exp_impl.ts b/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
new file mode 100644
index 00000000000..e8d8dd5b35e
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
@@ -0,0 +1,26 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {TypedArray} from '../../../types';
+
+export const exp = (x: TypedArray, outValues: TypedArray): TypedArray => {
+  for (let i = 0; i < x.length; ++i) {
+    outValues[i] = Math.exp(x[i]);
+  }
+
+  return outValues;
+};
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index 8b0f5f05457..7d9a8d4749b 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -21,6 +21,7 @@ import {parseAxisParam, sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
+import {exp} from './exp_impl';
 import {max} from './max_impl';
 import {sub} from './sub_impl';
 
@@ -64,6 +65,10 @@ registerKernel({
     console.log('subtract');
     console.log(a);
 
+    const b = exp(a, new Float32Array(sizeFromShape(logits.shape)));
+    console.log('exp');
+    console.log(b);
+
     const dataId = cpuBackend.write(maxLogit, maxLogitOutShape, logits.dtype);
     return {dataId, shape: maxLogitOutShape, dtype: logits.dtype};
   }

From 83e6d4606bba6dbb490691ca876ba7d4b70ed4c9 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 10:32:50 -0500
Subject: [PATCH 46/84] add sum

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  1 +
 tfjs-core/src/backends/cpu/kernels/Sum.ts     | 58 +++++++++++++++++++
 tfjs-core/src/backends/cpu/kernels/softmax.ts | 17 ++++--
 .../src/backends/cpu/kernels/sum_impl.ts      | 37 ++++++++++++
 4 files changed, 107 insertions(+), 6 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/Sum.ts
 create mode 100644 tfjs-core/src/backends/cpu/kernels/sum_impl.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index ef85b433af4..1b3a1d1d34e 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -24,3 +24,4 @@ import './kernels/softmax';
 import './kernels/max';
 import './kernels/Sub';
 import './kernels/Exp';
+import './kernels/Sum';
diff --git a/tfjs-core/src/backends/cpu/kernels/Sum.ts b/tfjs-core/src/backends/cpu/kernels/Sum.ts
new file mode 100644
index 00000000000..2b309ec077f
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/Sum.ts
@@ -0,0 +1,58 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. 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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {upcastType} from '../../../types';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+import {sum} from './sum_impl';
+
+interface SumInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface SumAttrs extends NamedAttrMap {
+  axes: number[];
+}
+
+registerKernel({
+  kernelName: 'Sum',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {x} = inputs as SumInputs;
+    const {axes} = attrs as SumAttrs;
+    const cpuBackend = backend as MathBackendCPU;
+
+    assertNotComplex(x, 'sum');
+
+    axis_util.assertAxesAreInnerMostDims('sum', axes, x.shape.length);
+
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(x.shape, axes);
+    const resultDtype = upcastType(x.dtype, 'int32');
+
+    const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
+    const result =
+        sum(xVals, reduceShape, new Float32Array(sizeFromShape(outShape)));
+
+    const dataId = cpuBackend.write(result, outShape, resultDtype);
+    return {dataId, shape: outShape, dtype: resultDtype};
+  }
+});
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index 7d9a8d4749b..cf0047c4b74 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -24,6 +24,7 @@ import {assertNotComplex} from '../cpu_util';
 import {exp} from './exp_impl';
 import {max} from './max_impl';
 import {sub} from './sub_impl';
+import {sum} from './sum_impl';
 
 interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
@@ -44,16 +45,15 @@ registerKernel({
 
     const axes = parseAxisParam([dim], logits.shape);
 
-    const [maxLogitOutShape, reduceShape] =
+    const [reduceOutShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
     const logitsValues =
         cpuBackend.data.get(logits.dataId).values as Float32Array;
     const maxLogit =
         max(logitsValues, reduceShape,
-            new Float32Array(sizeFromShape(maxLogitOutShape)));
+            new Float32Array(sizeFromShape(reduceOutShape)));
 
-    const expandedShape =
-        axis_util.expandShapeToKeepDim(maxLogitOutShape, axes);
+    const expandedShape = axis_util.expandShapeToKeepDim(reduceOutShape, axes);
 
     console.log('MAX LOGIT');
     console.log(maxLogit);
@@ -69,7 +69,12 @@ registerKernel({
     console.log('exp');
     console.log(b);
 
-    const dataId = cpuBackend.write(maxLogit, maxLogitOutShape, logits.dtype);
-    return {dataId, shape: maxLogitOutShape, dtype: logits.dtype};
+    const sumExp =
+        sum(b, reduceShape, new Float32Array(sizeFromShape(reduceOutShape)));
+    console.log('sumexp');
+    console.log(sumExp);
+
+    const dataId = cpuBackend.write(maxLogit, reduceOutShape, logits.dtype);
+    return {dataId, shape: reduceOutShape, dtype: logits.dtype};
   }
 });
diff --git a/tfjs-core/src/backends/cpu/kernels/sum_impl.ts b/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
new file mode 100644
index 00000000000..e0ba88d7295
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
@@ -0,0 +1,37 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {TypedArray} from '../../../types';
+import {sizeFromShape} from '../../../util';
+
+export const sum =
+    (x: TypedArray, reduceShape: number[], outValues: TypedArray):
+        TypedArray => {
+          const reduceSize = sizeFromShape(reduceShape);
+
+          for (let i = 0; i < x.length; ++i) {
+            const offset = i * reduceSize;
+            let sum = 0;
+            for (let j = 0; j < reduceSize; ++j) {
+              const value = x[offset + j];
+              sum += value;
+            }
+            outValues[i] = sum;
+          }
+
+          return outValues;
+        };

From 4e67f8d99a6dddff396b7adb55860ad61bdd796e Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 10:39:21 -0500
Subject: [PATCH 47/84] softmax

---
 tfjs-core/src/backends/cpu/all_kernels.ts     |  1 +
 tfjs-core/src/backends/cpu/kernels/Div.ts     | 51 ++++++++++++++++++
 .../src/backends/cpu/kernels/div_impl.ts      | 52 +++++++++++++++++++
 tfjs-core/src/backends/cpu/kernels/softmax.ts |  9 +++-
 4 files changed, 111 insertions(+), 2 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/Div.ts
 create mode 100644 tfjs-core/src/backends/cpu/kernels/div_impl.ts

diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
index 1b3a1d1d34e..13ee0f23aa5 100644
--- a/tfjs-core/src/backends/cpu/all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -25,3 +25,4 @@ import './kernels/max';
 import './kernels/Sub';
 import './kernels/Exp';
 import './kernels/Sum';
+import './kernels/Div';
diff --git a/tfjs-core/src/backends/cpu/kernels/Div.ts b/tfjs-core/src/backends/cpu/kernels/Div.ts
new file mode 100644
index 00000000000..46b48b87bc9
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/Div.ts
@@ -0,0 +1,51 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as broadcast_util from '../../../ops/broadcast_util';
+import {upcastType} from '../../../types';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+
+import {div} from './div_impl';
+
+interface DivInputs extends NamedTensorInfoMap {
+  a: TensorInfo;
+  b: TensorInfo;
+}
+
+registerKernel({
+  kernelName: 'Div',
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend}) => {
+    const {a, b} = inputs as DivInputs;
+    const cpuBackend = backend as MathBackendCPU;
+
+    const dtype = upcastType(a.dtype, b.dtype);
+    const outShape =
+        broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
+    const outValues = new Float32Array(sizeFromShape(outShape));
+
+    const aVals = cpuBackend.data.get(a.dataId).values as Float32Array;
+    const bVals = cpuBackend.data.get(b.dataId).values as Float32Array;
+
+    const result = div(aVals, a.shape, bVals, b.shape, outValues, outShape);
+
+    const dataId = cpuBackend.write(result, outShape, dtype);
+    return {dataId, shape: outShape, dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/cpu/kernels/div_impl.ts b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
new file mode 100644
index 00000000000..574dbdb348c
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 broadcast_util from '../../../ops/broadcast_util';
+import {TypedArray} from '../../../types';
+import * as util from '../../../util';
+
+export const div =
+    (a: TypedArray, aShape: number[], b: TypedArray, bShape: number[],
+     outValues: TypedArray, outShape: number[]): TypedArray => {
+      const aBroadcastDims = broadcast_util.getBroadcastDims(aShape, outShape);
+      const bBroadcastDims = broadcast_util.getBroadcastDims(bShape, outShape);
+
+      const aStrides = util.computeStrides(aShape);
+      const bStrides = util.computeStrides(bShape);
+      const outStrides = util.computeStrides(outShape);
+
+      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
+        for (let i = 0; i < outValues.length; ++i) {
+          outValues[i] = a[i % a.length] / b[i % b.length];
+        }
+      } else {
+        for (let i = 0; i < outValues.length; ++i) {
+          const loc = util.indexToLoc(i, outStrides);
+
+          const aLoc = loc.slice(-aShape.length);
+          aBroadcastDims.forEach(d => aLoc[d] = 0);
+          const aIndex = util.locToIndex(aLoc, aStrides);
+
+          const bLoc = loc.slice(-bShape.length);
+          bBroadcastDims.forEach(d => bLoc[d] = 0);
+          const bIndex = util.locToIndex(bLoc, bStrides);
+
+          outValues[i] = a[aIndex] / b[bIndex];
+        }
+      }
+      return outValues;
+    };
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index cf0047c4b74..9d330fa810d 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -21,6 +21,7 @@ import {parseAxisParam, sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
+import {div} from './div_impl';
 import {exp} from './exp_impl';
 import {max} from './max_impl';
 import {sub} from './sub_impl';
@@ -74,7 +75,11 @@ registerKernel({
     console.log('sumexp');
     console.log(sumExp);
 
-    const dataId = cpuBackend.write(maxLogit, reduceOutShape, logits.dtype);
-    return {dataId, shape: reduceOutShape, dtype: logits.dtype};
+    const out =
+        div(b, logits.shape, sumExp, expandedShape,
+            new Float32Array(sizeFromShape(logits.shape)), logits.shape);
+
+    const dataId = cpuBackend.write(out, logits.shape, logits.dtype);
+    return {dataId, shape: logits.shape, dtype: logits.dtype};
   }
 });

From 3bff003173c98f846cb4cd0211b66b41fa94022f Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 10:42:31 -0500
Subject: [PATCH 48/84] remove softmax

---
 tfjs-core/src/backends/cpu/backend_cpu.ts | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index af5cda78a23..ee2a1e6772a 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -377,18 +377,6 @@ export class MathBackendCPU extends KernelBackend {
     return result.toTensor() as T;
   }
 
-  // softmax<T extends Tensor>(logits: T, dim: number): T {
-  //   console.log('softmaxxxxx');
-  //   const axes = util.parseAxisParam([dim], logits.shape);
-  //   const maxLogit = this.max(logits, axes);
-  //   const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape,
-  //   axes); const a = this.subtract(logits, maxLogit.reshape(expandedShape));
-  //   const b = this.exp(a);
-  //   const sumExp = this.sum(b, axes).reshape(expandedShape);
-
-  //   return this.realDivide(b, sumExp) as T;
-  // }
-
   subtract(a: Tensor, b: Tensor): Tensor {
     if (a.dtype === 'complex64' || b.dtype === 'complex64') {
       return this.broadcastedBinaryComplexOp(

From 583a3f411a8fda3d453d056c86475c076ab14949 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 11:14:02 -0500
Subject: [PATCH 49/84] binary impl

---
 tfjs-core/src/backends/cpu/kernels/Sub.ts     | 34 +-----------
 .../src/backends/cpu/kernels/binary_impl.ts   | 52 +++++++++++++++++++
 .../src/backends/cpu/kernels/binary_kernel.ts | 51 ++++++++++++++++++
 .../src/backends/cpu/kernels/sub_impl.ts      | 36 +------------
 4 files changed, 107 insertions(+), 66 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/binary_impl.ts
 create mode 100644 tfjs-core/src/backends/cpu/kernels/binary_kernel.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/Sub.ts b/tfjs-core/src/backends/cpu/kernels/Sub.ts
index fa3bb841235..d82fab58a23 100644
--- a/tfjs-core/src/backends/cpu/kernels/Sub.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Sub.ts
@@ -15,37 +15,7 @@
  * =============================================================================
  */
 
-import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
-import * as broadcast_util from '../../../ops/broadcast_util';
-import {upcastType} from '../../../types';
-import {sizeFromShape} from '../../../util';
-import {MathBackendCPU} from '../backend_cpu';
-
+import {registerBinaryKernel} from './binary_kernel';
 import {sub} from './sub_impl';
 
-interface SubInputs extends NamedTensorInfoMap {
-  a: TensorInfo;
-  b: TensorInfo;
-}
-
-registerKernel({
-  kernelName: 'Sub',
-  backendName: 'cpu',
-  kernelFunc: ({inputs, backend}) => {
-    const {a, b} = inputs as SubInputs;
-    const cpuBackend = backend as MathBackendCPU;
-
-    const dtype = upcastType(a.dtype, b.dtype);
-    const outShape =
-        broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
-    const outValues = new Float32Array(sizeFromShape(outShape));
-
-    const aVals = cpuBackend.data.get(a.dataId).values as Float32Array;
-    const bVals = cpuBackend.data.get(b.dataId).values as Float32Array;
-
-    const result = sub(aVals, a.shape, bVals, b.shape, outValues, outShape);
-
-    const dataId = cpuBackend.write(result, outShape, dtype);
-    return {dataId, shape: outShape, dtype};
-  }
-});
+registerBinaryKernel('Sub', sub);
diff --git a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
new file mode 100644
index 00000000000..d044bedee81
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 broadcast_util from '../../../ops/broadcast_util';
+import {TypedArray} from '../../../types';
+import * as util from '../../../util';
+
+export const createBinaryOp = (op: any) =>
+    (a: TypedArray, aShape: number[], b: TypedArray, bShape: number[],
+     outValues: TypedArray, outShape: number[]): TypedArray => {
+      const aBroadcastDims = broadcast_util.getBroadcastDims(aShape, outShape);
+      const bBroadcastDims = broadcast_util.getBroadcastDims(bShape, outShape);
+
+      const aStrides = util.computeStrides(aShape);
+      const bStrides = util.computeStrides(bShape);
+      const outStrides = util.computeStrides(outShape);
+
+      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
+        for (let i = 0; i < outValues.length; ++i) {
+          outValues[i] = op(a[i % a.length], b[i % b.length]);
+        }
+      } else {
+        for (let i = 0; i < outValues.length; ++i) {
+          const loc = util.indexToLoc(i, outStrides);
+
+          const aLoc = loc.slice(-aShape.length);
+          aBroadcastDims.forEach(d => aLoc[d] = 0);
+          const aIndex = util.locToIndex(aLoc, aStrides);
+
+          const bLoc = loc.slice(-bShape.length);
+          bBroadcastDims.forEach(d => bLoc[d] = 0);
+          const bIndex = util.locToIndex(bLoc, bStrides);
+
+          outValues[i] = op(a[aIndex], b[bIndex]);
+        }
+      }
+      return outValues;
+    };
diff --git a/tfjs-core/src/backends/cpu/kernels/binary_kernel.ts b/tfjs-core/src/backends/cpu/kernels/binary_kernel.ts
new file mode 100644
index 00000000000..b7e528e3c73
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/binary_kernel.ts
@@ -0,0 +1,51 @@
+/**
+ * @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 {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as broadcast_util from '../../../ops/broadcast_util';
+import {upcastType} from '../../../types';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+
+export function registerBinaryKernel(kernelName: string, op: any) {
+  registerKernel({
+    kernelName,
+    backendName: 'cpu',
+    kernelFunc: ({inputs, backend}) => {
+      const {a, b} = inputs as BinaryInputs;
+      const cpuBackend = backend as MathBackendCPU;
+
+      const dtype = upcastType(a.dtype, b.dtype);
+      const outShape =
+          broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
+      const outValues = new Float32Array(sizeFromShape(outShape));
+
+      const aVals = cpuBackend.data.get(a.dataId).values as Float32Array;
+      const bVals = cpuBackend.data.get(b.dataId).values as Float32Array;
+
+      const result = op(aVals, a.shape, bVals, b.shape, outValues, outShape);
+
+      const dataId = cpuBackend.write(result, outShape, dtype);
+      return {dataId, shape: outShape, dtype};
+    }
+  });
+}
+
+interface BinaryInputs extends NamedTensorInfoMap {
+  a: TensorInfo;
+  b: TensorInfo;
+}
diff --git a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
index a6a6b34335d..749e59da7e7 100644
--- a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
@@ -15,38 +15,6 @@
  * =============================================================================
  */
 
-import * as broadcast_util from '../../../ops/broadcast_util';
-import {TypedArray} from '../../../types';
-import * as util from '../../../util';
+import {createBinaryOp} from './binary_impl';
 
-export const sub =
-    (a: TypedArray, aShape: number[], b: TypedArray, bShape: number[],
-     outValues: TypedArray, outShape: number[]): TypedArray => {
-      const aBroadcastDims = broadcast_util.getBroadcastDims(aShape, outShape);
-      const bBroadcastDims = broadcast_util.getBroadcastDims(bShape, outShape);
-
-      const aStrides = util.computeStrides(aShape);
-      const bStrides = util.computeStrides(bShape);
-      const outStrides = util.computeStrides(outShape);
-
-      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
-        for (let i = 0; i < outValues.length; ++i) {
-          outValues[i] = a[i % a.length] - b[i % b.length];
-        }
-      } else {
-        for (let i = 0; i < outValues.length; ++i) {
-          const loc = util.indexToLoc(i, outStrides);
-
-          const aLoc = loc.slice(-aShape.length);
-          aBroadcastDims.forEach(d => aLoc[d] = 0);
-          const aIndex = util.locToIndex(aLoc, aStrides);
-
-          const bLoc = loc.slice(-bShape.length);
-          bBroadcastDims.forEach(d => bLoc[d] = 0);
-          const bIndex = util.locToIndex(bLoc, bStrides);
-
-          outValues[i] = a[aIndex] - b[bIndex];
-        }
-      }
-      return outValues;
-    };
+export const sub = createBinaryOp((a: number, b: number) => a - b);

From cb5c1a574626d500c8e25a05f001faebd04e7042 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 14 Feb 2020 11:15:44 -0500
Subject: [PATCH 50/84] clean

---
 tfjs-core/src/backends/cpu/kernels/Div.ts     | 34 ++----------------
 .../src/backends/cpu/kernels/div_impl.ts      | 36 ++-----------------
 2 files changed, 4 insertions(+), 66 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/kernels/Div.ts b/tfjs-core/src/backends/cpu/kernels/Div.ts
index 46b48b87bc9..3efc0ed7453 100644
--- a/tfjs-core/src/backends/cpu/kernels/Div.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Div.ts
@@ -15,37 +15,7 @@
  * =============================================================================
  */
 
-import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
-import * as broadcast_util from '../../../ops/broadcast_util';
-import {upcastType} from '../../../types';
-import {sizeFromShape} from '../../../util';
-import {MathBackendCPU} from '../backend_cpu';
-
+import {registerBinaryKernel} from './binary_kernel';
 import {div} from './div_impl';
 
-interface DivInputs extends NamedTensorInfoMap {
-  a: TensorInfo;
-  b: TensorInfo;
-}
-
-registerKernel({
-  kernelName: 'Div',
-  backendName: 'cpu',
-  kernelFunc: ({inputs, backend}) => {
-    const {a, b} = inputs as DivInputs;
-    const cpuBackend = backend as MathBackendCPU;
-
-    const dtype = upcastType(a.dtype, b.dtype);
-    const outShape =
-        broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
-    const outValues = new Float32Array(sizeFromShape(outShape));
-
-    const aVals = cpuBackend.data.get(a.dataId).values as Float32Array;
-    const bVals = cpuBackend.data.get(b.dataId).values as Float32Array;
-
-    const result = div(aVals, a.shape, bVals, b.shape, outValues, outShape);
-
-    const dataId = cpuBackend.write(result, outShape, dtype);
-    return {dataId, shape: outShape, dtype};
-  }
-});
+registerBinaryKernel('Div', div);
diff --git a/tfjs-core/src/backends/cpu/kernels/div_impl.ts b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
index 574dbdb348c..ed14151713f 100644
--- a/tfjs-core/src/backends/cpu/kernels/div_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
@@ -15,38 +15,6 @@
  * =============================================================================
  */
 
-import * as broadcast_util from '../../../ops/broadcast_util';
-import {TypedArray} from '../../../types';
-import * as util from '../../../util';
+import {createBinaryOp} from './binary_impl';
 
-export const div =
-    (a: TypedArray, aShape: number[], b: TypedArray, bShape: number[],
-     outValues: TypedArray, outShape: number[]): TypedArray => {
-      const aBroadcastDims = broadcast_util.getBroadcastDims(aShape, outShape);
-      const bBroadcastDims = broadcast_util.getBroadcastDims(bShape, outShape);
-
-      const aStrides = util.computeStrides(aShape);
-      const bStrides = util.computeStrides(bShape);
-      const outStrides = util.computeStrides(outShape);
-
-      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
-        for (let i = 0; i < outValues.length; ++i) {
-          outValues[i] = a[i % a.length] / b[i % b.length];
-        }
-      } else {
-        for (let i = 0; i < outValues.length; ++i) {
-          const loc = util.indexToLoc(i, outStrides);
-
-          const aLoc = loc.slice(-aShape.length);
-          aBroadcastDims.forEach(d => aLoc[d] = 0);
-          const aIndex = util.locToIndex(aLoc, aStrides);
-
-          const bLoc = loc.slice(-bShape.length);
-          bBroadcastDims.forEach(d => bLoc[d] = 0);
-          const bIndex = util.locToIndex(bLoc, bStrides);
-
-          outValues[i] = a[aIndex] / b[bIndex];
-        }
-      }
-      return outValues;
-    };
+export const div = createBinaryOp((a: number, b: number) => a / b);

From a2a5792182e28be014c884d89c436186c49d9d54 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Mon, 17 Feb 2020 11:43:13 -0500
Subject: [PATCH 51/84] webgl wip

---
 tfjs-core/src/backends/webgl/all_kernels.ts   |  1 +
 tfjs-core/src/backends/webgl/backend_webgl.ts | 20 ++---
 .../src/backends/webgl/kernels/Softmax.ts     | 75 +++++++++++++++++++
 tfjs-core/src/ops/softmax_test.ts             |  5 +-
 4 files changed, 89 insertions(+), 12 deletions(-)
 create mode 100644 tfjs-core/src/backends/webgl/kernels/Softmax.ts

diff --git a/tfjs-core/src/backends/webgl/all_kernels.ts b/tfjs-core/src/backends/webgl/all_kernels.ts
index 724e83ce034..7f9b3fa9039 100644
--- a/tfjs-core/src/backends/webgl/all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/all_kernels.ts
@@ -19,5 +19,6 @@
 // global registry when we compile the library. A modular build would replace
 // the contents of this file and import only the kernels that are needed.
 import './square';
+import './kernels/Softmax';
 import './fromPixels';
 import './non_max_suppression_v5';
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index a1dfe5c0086..fbeb3734763 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1597,16 +1597,16 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [x]);
   }
 
-  softmax<T extends Tensor>(logits: T, dim: number): T {
-    const axes = util.parseAxisParam([dim], logits.shape);
-    const maxLogit = this.max(logits, axes);
-    const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape, axes);
-    const a = this.subtract(logits, maxLogit.reshape(expandedShape));
-    const b = this.exp(a);
-    const sumExp = this.sum(b, axes).reshape(expandedShape);
+  // softmax<T extends Tensor>(logits: T, dim: number): T {
+  //   const axes = util.parseAxisParam([dim], logits.shape);
+  //   const maxLogit = this.max(logits, axes);
+  //   const expandedShape = axis_util.expandShapeToKeepDim(maxLogit.shape,
+  //   axes); const a = this.subtract(logits, maxLogit.reshape(expandedShape));
+  //   const b = this.exp(a);
+  //   const sumExp = this.sum(b, axes).reshape(expandedShape);
 
-    return this.realDivide(b, sumExp) as T;
-  }
+  //   return this.realDivide(b, sumExp) as T;
+  // }
 
   log<T extends Tensor>(x: T): T {
     if (this.shouldExecuteOnCPU([x])) {
@@ -2785,7 +2785,7 @@ export class MathBackendWebGL extends KernelBackend {
     return texData.values as TypedArray;
   }
 
-  private acquireTexture(
+  acquireTexture(
       texShape: [number, number], texType: TextureUsage, dtype: DataType,
       isPacked: boolean): WebGLTexture {
     this.numBytesInGPU += this.computeBytes(texShape, dtype);
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
new file mode 100644
index 00000000000..3239fe067c8
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -0,0 +1,75 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {parseAxisParam, sizeFromShape} from '../../../util';
+import {webgl_util} from '../../../webgl';
+import {MathBackendWebGL} from '../backend_webgl';
+
+interface SoftmaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface SoftmaxAttrs extends NamedAttrMap {
+  dim: number;
+}
+
+registerKernel({
+  kernelName: 'Softmax',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {logits} = inputs as SoftmaxInputs;
+    const {dim} = attrs as SoftmaxAttrs;
+    const webglBackend = backend as MathBackendWebGL;
+    const textureManager = webglBackend.getTextureManager();
+
+    const axes = parseAxisParam([dim], logits.shape);
+
+    const [reduceOutShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(logits.shape, axes);
+    // const logitsValues =
+    //     webglBackend.data.get(logits.dataId).values as Float32Array;
+    // const maxLogit =
+    //     max(logitsValues, reduceShape,
+    //         new Float32Array(sizeFromShape(reduceOutShape)));
+
+
+    // Create resource for max kernel:
+    const logitsTexdata = webglBackend.texData.get(logits.dataId);
+
+    const texShapeForMax =
+        webgl_util.getTextureShapeFromLogicalShape(reduceOutShape);
+
+    // Not sure about usage - check webglBackend.runWebglProgram for reference
+    // on shader program output usage.
+    const texUsage: any = null;
+    // This is only the texture... need to create texData object.
+    const texForMax = webglBackend.acquireTexture(
+        texShapeForMax, texUsage, logits.dtype, logitsTexdata.isPacked);
+
+    const maxLogit = max(
+        logitsTexdata,
+        reduceShape,
+    );
+
+    console.log('MAX');
+
+    const dataId = webglBackend.write(maxLogit, logits.shape, logits.dtype);
+    return {dataId, shape: logits.shape, dtype: logits.dtype};
+  }
+});
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 9915fca14a7..78a91462189 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,12 +15,13 @@
  * =============================================================================
  */
 
-import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
+// import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
+import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', CPU_ENVS, () => {
+describeWithFlags('softmax', WEBGL_ENVS, () => {
   fit('regular test', async () => {
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From 93d8c4130401aef6b1dcebf2b56b5d1f7d2487ff Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 10:50:31 -0500
Subject: [PATCH 52/84] max webgl

---
 tfjs-core/src/backends/webgl/all_kernels.ts   |  1 +
 tfjs-core/src/backends/webgl/backend_webgl.ts | 24 +++----
 tfjs-core/src/backends/webgl/kernels/Max.ts   | 69 +++++++++++++++++++
 .../src/backends/webgl/kernels/Softmax.ts     | 17 +++--
 tfjs-core/src/ops/softmax_test.ts             |  9 ++-
 5 files changed, 98 insertions(+), 22 deletions(-)
 create mode 100644 tfjs-core/src/backends/webgl/kernels/Max.ts

diff --git a/tfjs-core/src/backends/webgl/all_kernels.ts b/tfjs-core/src/backends/webgl/all_kernels.ts
index 7f9b3fa9039..9b90c8a1967 100644
--- a/tfjs-core/src/backends/webgl/all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/all_kernels.ts
@@ -20,5 +20,6 @@
 // the contents of this file and import only the kernels that are needed.
 import './square';
 import './kernels/Softmax';
+import './kernels/Max';
 import './fromPixels';
 import './non_max_suppression_v5';
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index fbeb3734763..f367284a2a7 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1330,18 +1330,18 @@ export class MathBackendWebGL extends KernelBackend {
     return this.compileAndRun(program, [a, b]);
   }
 
-  max(x: Tensor, axes: number[]): Tensor {
-    if (this.shouldExecuteOnCPU([x])) {
-      return this.cpuBackend.max(x, axes);
-    }
-
-    axis_util.assertAxesAreInnerMostDims('max', axes, x.rank);
-    const [outShape, reduceShape] =
-        axis_util.computeOutAndReduceShapes(x.shape, axes);
-    const inSize = util.sizeFromShape(reduceShape);
-    const a2D = x.as2D(-1, inSize);
-    return this.reduce(a2D, 'max', a2D.dtype).reshape(outShape);
-  }
+  // max(x: Tensor, axes: number[]): Tensor {
+  //   if (this.shouldExecuteOnCPU([x])) {
+  //     return this.cpuBackend.max(x, axes);
+  //   }
+
+  //   axis_util.assertAxesAreInnerMostDims('max', axes, x.rank);
+  //   const [outShape, reduceShape] =
+  //       axis_util.computeOutAndReduceShapes(x.shape, axes);
+  //   const inSize = util.sizeFromShape(reduceShape);
+  //   const a2D = x.as2D(-1, inSize);
+  //   return this.reduce(a2D, 'max', a2D.dtype).reshape(outShape);
+  // }
 
   maximum(a: Tensor, b: Tensor): Tensor {
     if (this.shouldExecuteOnCPU([a, b])) {
diff --git a/tfjs-core/src/backends/webgl/kernels/Max.ts b/tfjs-core/src/backends/webgl/kernels/Max.ts
new file mode 100644
index 00000000000..e4dbd4f1586
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/kernels/Max.ts
@@ -0,0 +1,69 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {computeOptimalWindowSize} from '../../../ops/reduce_util';
+import {sizeFromShape} from '../../../util';
+import {MathBackendWebGL} from '../backend_webgl';
+import {ReduceProgram} from '../reduce_gpu';
+
+interface MaxInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface MaxAttrs extends NamedAttrMap {
+  axes: number[];
+}
+
+export const maxImpl =
+    (x: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
+      const [batchSize, inSize] = x.shape;
+      const windowSize = computeOptimalWindowSize(inSize);
+      const reduceInfo = {windowSize, inSize, batchSize};
+      const program = new ReduceProgram(reduceInfo, 'max');
+      const output = backend.runWebGLProgram(program, [x], x.dtype);
+
+      if (output.shape[1] === 1) {
+        return output;
+      }
+      return maxImpl(output, backend);
+    };
+
+registerKernel({
+  kernelName: 'Max',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {x} = inputs as MaxInputs;
+    const {axes} = attrs as MaxAttrs;
+    const webglBackend = backend as MathBackendWebGL;
+
+    axis_util.assertAxesAreInnerMostDims('max', axes, x.shape.length);
+
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(x.shape, axes);
+    const inSize = sizeFromShape(reduceShape);
+    const xSize = sizeFromShape(x.shape);
+
+    // TODO: Call reshape kernel.
+    x.shape = [xSize / inSize, inSize];
+
+    const out = maxImpl(x, webglBackend);
+
+    return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index 3239fe067c8..93eaa2b0b5c 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -36,7 +36,7 @@ registerKernel({
     const {logits} = inputs as SoftmaxInputs;
     const {dim} = attrs as SoftmaxAttrs;
     const webglBackend = backend as MathBackendWebGL;
-    const textureManager = webglBackend.getTextureManager();
+    // const textureManager = webglBackend.getTextureManager();
 
     const axes = parseAxisParam([dim], logits.shape);
 
@@ -62,14 +62,17 @@ registerKernel({
     const texForMax = webglBackend.acquireTexture(
         texShapeForMax, texUsage, logits.dtype, logitsTexdata.isPacked);
 
-    const maxLogit = max(
-        logitsTexdata,
-        reduceShape,
-    );
+    // const maxLogit = max(
+    //     logitsTexdata,
+    //     reduceShape,
+    // );
 
     console.log('MAX');
+    console.log(texForMax, reduceShape, sizeFromShape);
 
-    const dataId = webglBackend.write(maxLogit, logits.shape, logits.dtype);
-    return {dataId, shape: logits.shape, dtype: logits.dtype};
+    return {dataId: logits.dataId, shape: logits.shape, dtype: logits.dtype};
+
+    // const dataId = webglBackend.write(maxLogit, logits.shape, logits.dtype);
+    // return {dataId, shape: logits.shape, dtype: logits.dtype};
   }
 });
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 78a91462189..401ef48cff3 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -23,10 +23,13 @@ import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
   fit('regular test', async () => {
-    const y = tf.softmax(tf.tensor1d([2, 1, 3]));
+    console.log('TESTINGGGG');
+    // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
-    expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
-    expectArraysClose(await y.sum().data(), 1);
+    // expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
+    // expectArraysClose(await y.sum().data(), 1);
+    const x = tf.tensor1d([1, 2, 3]);
+    console.log(x.max());
   });
 });
 

From 277d2468a3805fac5c4d47ed17e294d53df9929b Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 11:00:25 -0500
Subject: [PATCH 53/84] clean

---
 tfjs-core/src/backends/webgl/kernels/Max.ts   | 36 ++++++++--------
 .../src/backends/webgl/kernels/Softmax.ts     | 41 ++++---------------
 2 files changed, 26 insertions(+), 51 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Max.ts b/tfjs-core/src/backends/webgl/kernels/Max.ts
index e4dbd4f1586..b591206cc00 100644
--- a/tfjs-core/src/backends/webgl/kernels/Max.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Max.ts
@@ -31,18 +31,25 @@ interface MaxAttrs extends NamedAttrMap {
 }
 
 export const maxImpl =
-    (x: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
-      const [batchSize, inSize] = x.shape;
-      const windowSize = computeOptimalWindowSize(inSize);
-      const reduceInfo = {windowSize, inSize, batchSize};
-      const program = new ReduceProgram(reduceInfo, 'max');
-      const output = backend.runWebGLProgram(program, [x], x.dtype);
+    (x: TensorInfo, reduceShape: number[], backend: MathBackendWebGL):
+        TensorInfo => {
+          const inSize = sizeFromShape(reduceShape);
+          const xSize = sizeFromShape(x.shape);
+          const batchSize = xSize / inSize;
 
-      if (output.shape[1] === 1) {
-        return output;
-      }
-      return maxImpl(output, backend);
-    };
+          // TODO: Call reshape kernel.
+          x.shape = [batchSize, inSize];
+
+          const windowSize = computeOptimalWindowSize(inSize);
+          const reduceInfo = {windowSize, inSize, batchSize};
+          const program = new ReduceProgram(reduceInfo, 'max');
+          const output = backend.runWebGLProgram(program, [x], x.dtype);
+
+          if (output.shape[1] === 1) {
+            return output;
+          }
+          return maxImpl(output, reduceShape, backend);
+        };
 
 registerKernel({
   kernelName: 'Max',
@@ -56,13 +63,8 @@ registerKernel({
 
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(x.shape, axes);
-    const inSize = sizeFromShape(reduceShape);
-    const xSize = sizeFromShape(x.shape);
-
-    // TODO: Call reshape kernel.
-    x.shape = [xSize / inSize, inSize];
 
-    const out = maxImpl(x, webglBackend);
+    const out = maxImpl(x, reduceShape, webglBackend);
 
     return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
   }
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index 93eaa2b0b5c..292eb65b9d1 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -17,9 +17,10 @@
 
 import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
-import {parseAxisParam, sizeFromShape} from '../../../util';
-import {webgl_util} from '../../../webgl';
+// import {parseAxisParam, sizeFromShape} from '../../../util';
+import {parseAxisParam} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
+import {maxImpl} from './Max';
 
 interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
@@ -36,43 +37,15 @@ registerKernel({
     const {logits} = inputs as SoftmaxInputs;
     const {dim} = attrs as SoftmaxAttrs;
     const webglBackend = backend as MathBackendWebGL;
-    // const textureManager = webglBackend.getTextureManager();
 
     const axes = parseAxisParam([dim], logits.shape);
 
-    const [reduceOutShape, reduceShape] =
+    const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
-    // const logitsValues =
-    //     webglBackend.data.get(logits.dataId).values as Float32Array;
-    // const maxLogit =
-    //     max(logitsValues, reduceShape,
-    //         new Float32Array(sizeFromShape(reduceOutShape)));
 
+    const out = maxImpl(logits, reduceShape, webglBackend);
+    console.log(outShape);
 
-    // Create resource for max kernel:
-    const logitsTexdata = webglBackend.texData.get(logits.dataId);
-
-    const texShapeForMax =
-        webgl_util.getTextureShapeFromLogicalShape(reduceOutShape);
-
-    // Not sure about usage - check webglBackend.runWebglProgram for reference
-    // on shader program output usage.
-    const texUsage: any = null;
-    // This is only the texture... need to create texData object.
-    const texForMax = webglBackend.acquireTexture(
-        texShapeForMax, texUsage, logits.dtype, logitsTexdata.isPacked);
-
-    // const maxLogit = max(
-    //     logitsTexdata,
-    //     reduceShape,
-    // );
-
-    console.log('MAX');
-    console.log(texForMax, reduceShape, sizeFromShape);
-
-    return {dataId: logits.dataId, shape: logits.shape, dtype: logits.dtype};
-
-    // const dataId = webglBackend.write(maxLogit, logits.shape, logits.dtype);
-    // return {dataId, shape: logits.shape, dtype: logits.dtype};
+    return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
   }
 });

From 66cabfb4edbae52fa1720d611f0b1b6cd0245914 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 12:18:04 -0500
Subject: [PATCH 54/84] add reshape

---
 tfjs-core/src/backends/webgl/kernels/Max.ts |  4 +-
 tfjs-core/src/backends/webgl/reshape.ts     | 65 +++++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)
 create mode 100644 tfjs-core/src/backends/webgl/reshape.ts

diff --git a/tfjs-core/src/backends/webgl/kernels/Max.ts b/tfjs-core/src/backends/webgl/kernels/Max.ts
index b591206cc00..fb4bd83fc19 100644
--- a/tfjs-core/src/backends/webgl/kernels/Max.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Max.ts
@@ -21,6 +21,7 @@ import {computeOptimalWindowSize} from '../../../ops/reduce_util';
 import {sizeFromShape} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
 import {ReduceProgram} from '../reduce_gpu';
+import {reshape} from '../reshape';
 
 interface MaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
@@ -37,8 +38,7 @@ export const maxImpl =
           const xSize = sizeFromShape(x.shape);
           const batchSize = xSize / inSize;
 
-          // TODO: Call reshape kernel.
-          x.shape = [batchSize, inSize];
+          x = reshape(x, [batchSize, inSize], backend);
 
           const windowSize = computeOptimalWindowSize(inSize);
           const reduceInfo = {windowSize, inSize, batchSize};
diff --git a/tfjs-core/src/backends/webgl/reshape.ts b/tfjs-core/src/backends/webgl/reshape.ts
new file mode 100644
index 00000000000..00d187b092c
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/reshape.ts
@@ -0,0 +1,65 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {TensorInfo} from '../../kernel_registry';
+import {webgl_util} from '../../webgl';
+
+import {MathBackendWebGL} from './backend_webgl';
+import {ReshapePackedProgram} from './reshape_packed_gpu';
+
+const packedReshape =
+    (input: TensorInfo, afterShape: number[], backend: MathBackendWebGL):
+        TensorInfo => {
+          const input3DShape = [
+            webgl_util.getBatchDim(input.shape),
+            ...webgl_util.getRowsCols(input.shape)
+          ] as [number, number, number];
+          const input3D: TensorInfo = {
+            dtype: input.dtype,
+            shape: input3DShape,
+            dataId: input.dataId
+          };
+          const afterShapeAs3D = [
+            webgl_util.getBatchDim(afterShape),
+            ...webgl_util.getRowsCols(afterShape)
+          ] as [number, number, number];
+
+          const program =
+              new ReshapePackedProgram(afterShapeAs3D, input3DShape);
+          const preventEagerUnpackingOfOutput = true;
+          const output = backend.runWebGLProgram(
+              program, [input3D], input.dtype, null /* customSetup */,
+              preventEagerUnpackingOfOutput);
+          return {
+            dataId: output.dataId,
+            shape: afterShape,
+            dtype: output.dtype
+          };
+        }
+
+export const reshape =
+    (x: TensorInfo, afterShape: number[],
+     backend: MathBackendWebGL): TensorInfo => {
+      const xTexData = backend.texData.get(x.dataId);
+      if (xTexData.isPacked && !webgl_util.isReshapeFree(x.shape, afterShape) &&
+          !(xTexData.texture !== null &&
+            webgl_util.isReshapeFree(xTexData.shape, afterShape))) {
+        return packedReshape(x, afterShape, backend);
+      }
+      x.shape = afterShape;
+      return x;
+    };

From 94eb1517fd0ec9bde230d4eb16c737651248e96a Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 13:26:59 -0500
Subject: [PATCH 55/84] add sub shell

---
 tfjs-core/src/backends/webgl/all_kernels.ts |  1 +
 tfjs-core/src/backends/webgl/kernels/Sub.ts | 52 +++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 tfjs-core/src/backends/webgl/kernels/Sub.ts

diff --git a/tfjs-core/src/backends/webgl/all_kernels.ts b/tfjs-core/src/backends/webgl/all_kernels.ts
index 9b90c8a1967..b2a2903cc1d 100644
--- a/tfjs-core/src/backends/webgl/all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/all_kernels.ts
@@ -21,5 +21,6 @@
 import './square';
 import './kernels/Softmax';
 import './kernels/Max';
+import './kernels/Sub';
 import './fromPixels';
 import './non_max_suppression_v5';
diff --git a/tfjs-core/src/backends/webgl/kernels/Sub.ts b/tfjs-core/src/backends/webgl/kernels/Sub.ts
new file mode 100644
index 00000000000..6f487d0a5ff
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/kernels/Sub.ts
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {env} from '../../../environment';
+import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {upcastType} from '../../../types';
+import {MathBackendWebGL} from '../backend_webgl';
+import * as binaryop_gpu from '../binaryop_gpu';
+import {BinaryOpProgram} from '../binaryop_gpu';
+import {BinaryOpPackedProgram} from '../binaryop_packed_gpu';
+
+interface SubInputs extends NamedTensorInfoMap {
+  a: TensorInfo;
+  b: TensorInfo;
+}
+
+export const subImpl =
+    (a: TensorInfo, b: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
+      const dtype = upcastType(a.dtype, b.dtype);
+      let program = new BinaryOpProgram(binaryop_gpu.SUB, a.shape, b.shape);
+      if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
+        program = new BinaryOpPackedProgram(binaryop_gpu.SUB, a.shape, b.shape);
+      }
+      const output = backend.runWebGLProgram(program, [a, b], dtype);
+      return output;
+    };
+
+registerKernel({
+  kernelName: 'Sub',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, backend}) => {
+    const {a, b} = inputs as SubInputs;
+    const webglBackend = backend as MathBackendWebGL;
+    const out = subImpl(a, b, webglBackend);
+
+    return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
+  }
+})

From 1716084eae11fb4ba32a03f0ab377b11e0108799 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 18:15:59 -0500
Subject: [PATCH 56/84] add sub

---
 tfjs-core/src/backends/webgl/kernels/Softmax.ts | 9 ++++++++-
 tfjs-core/src/backends/webgl/kernels/Sub.ts     | 2 +-
 tfjs-core/src/ops/softmax_test.ts               | 2 +-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index 292eb65b9d1..f13cc26b599 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -21,6 +21,7 @@ import * as axis_util from '../../../ops/axis_util';
 import {parseAxisParam} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
 import {maxImpl} from './Max';
+import {subImpl} from './Sub';
 
 interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
@@ -43,9 +44,15 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
 
-    const out = maxImpl(logits, reduceShape, webglBackend);
+    const max = maxImpl(logits, reduceShape, webglBackend);
     console.log(outShape);
 
+    const out = subImpl(logits, max, webglBackend);
+
+    webglBackend.disposeData(max.dataId);
+
+    console.log('RAN SUB');
+
     return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
   }
 });
diff --git a/tfjs-core/src/backends/webgl/kernels/Sub.ts b/tfjs-core/src/backends/webgl/kernels/Sub.ts
index 6f487d0a5ff..5f017a353a0 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sub.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sub.ts
@@ -49,4 +49,4 @@ registerKernel({
 
     return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
   }
-})
+});
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 401ef48cff3..43260409c4b 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -29,7 +29,7 @@ describeWithFlags('softmax', WEBGL_ENVS, () => {
     // expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
     // expectArraysClose(await y.sum().data(), 1);
     const x = tf.tensor1d([1, 2, 3]);
-    console.log(x.max());
+    x.softmax().print();
   });
 });
 

From ab272bd807c54401af9adfa3fc941c0033790a00 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 18:20:49 -0500
Subject: [PATCH 57/84] add exp

---
 tfjs-core/src/backends/webgl/all_kernels.ts   |  1 +
 tfjs-core/src/backends/webgl/kernels/Exp.ts   | 49 +++++++++++++++++++
 .../src/backends/webgl/kernels/Softmax.ts     |  6 ++-
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 tfjs-core/src/backends/webgl/kernels/Exp.ts

diff --git a/tfjs-core/src/backends/webgl/all_kernels.ts b/tfjs-core/src/backends/webgl/all_kernels.ts
index b2a2903cc1d..300ff8e7a53 100644
--- a/tfjs-core/src/backends/webgl/all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/all_kernels.ts
@@ -22,5 +22,6 @@ import './square';
 import './kernels/Softmax';
 import './kernels/Max';
 import './kernels/Sub';
+import './kernels/Exp';
 import './fromPixels';
 import './non_max_suppression_v5';
diff --git a/tfjs-core/src/backends/webgl/kernels/Exp.ts b/tfjs-core/src/backends/webgl/kernels/Exp.ts
new file mode 100644
index 00000000000..de8fd339994
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/kernels/Exp.ts
@@ -0,0 +1,49 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {env} from '../../../environment';
+import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {MathBackendWebGL} from '../backend_webgl';
+import * as unaryop_gpu from '../unaryop_gpu';
+import {UnaryOpProgram} from '../unaryop_gpu';
+import {UnaryOpPackedProgram} from '../unaryop_packed_gpu';
+
+interface ExpInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+export const expImpl =
+    (x: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
+      let program = new UnaryOpProgram(x.shape, unaryop_gpu.EXP);
+      if (env().getBool('WEBGL_PACK_UNARY_OPERATIONS')) {
+        program = new UnaryOpPackedProgram(x.shape, unaryop_gpu.EXP);
+      }
+      const output = backend.runWebGLProgram(program, [x], x.dtype);
+      return output;
+    };
+
+registerKernel({
+  kernelName: 'Exp',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, backend}) => {
+    const {x} = inputs as ExpInputs;
+    const webglBackend = backend as MathBackendWebGL;
+    const out = expImpl(x, webglBackend);
+
+    return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index f13cc26b599..b1d90fe85dd 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -20,6 +20,8 @@ import * as axis_util from '../../../ops/axis_util';
 // import {parseAxisParam, sizeFromShape} from '../../../util';
 import {parseAxisParam} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
+
+import {expImpl} from './Exp';
 import {maxImpl} from './Max';
 import {subImpl} from './Sub';
 
@@ -47,9 +49,11 @@ registerKernel({
     const max = maxImpl(logits, reduceShape, webglBackend);
     console.log(outShape);
 
-    const out = subImpl(logits, max, webglBackend);
+    const subtracted = subImpl(logits, max, webglBackend);
+    const out = expImpl(subtracted, webglBackend);
 
     webglBackend.disposeData(max.dataId);
+    webglBackend.disposeData(subtracted.dataId);
 
     console.log('RAN SUB');
 

From bf3af8ec14af12967a2bd8f89e9cfee0c861c480 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Tue, 18 Feb 2020 18:25:41 -0500
Subject: [PATCH 58/84] add sum

---
 tfjs-core/src/backends/webgl/all_kernels.ts   |  1 +
 .../src/backends/webgl/kernels/Softmax.ts     |  6 +-
 tfjs-core/src/backends/webgl/kernels/Sum.ts   | 72 +++++++++++++++++++
 3 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 tfjs-core/src/backends/webgl/kernels/Sum.ts

diff --git a/tfjs-core/src/backends/webgl/all_kernels.ts b/tfjs-core/src/backends/webgl/all_kernels.ts
index 300ff8e7a53..7ed744f8b5f 100644
--- a/tfjs-core/src/backends/webgl/all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/all_kernels.ts
@@ -21,6 +21,7 @@
 import './square';
 import './kernels/Softmax';
 import './kernels/Max';
+import './kernels/Sum';
 import './kernels/Sub';
 import './kernels/Exp';
 import './fromPixels';
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index b1d90fe85dd..e891f0b6d67 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -24,6 +24,7 @@ import {MathBackendWebGL} from '../backend_webgl';
 import {expImpl} from './Exp';
 import {maxImpl} from './Max';
 import {subImpl} from './Sub';
+import {sumImpl} from './Sum';
 
 interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
@@ -50,10 +51,13 @@ registerKernel({
     console.log(outShape);
 
     const subtracted = subImpl(logits, max, webglBackend);
-    const out = expImpl(subtracted, webglBackend);
+    const exponentiated = expImpl(subtracted, webglBackend);
+
+    const out = sumImpl(exponentiated, reduceShape, webglBackend);
 
     webglBackend.disposeData(max.dataId);
     webglBackend.disposeData(subtracted.dataId);
+    webglBackend.disposeData(exponentiated.dataId);
 
     console.log('RAN SUB');
 
diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
new file mode 100644
index 00000000000..65609649ac7
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -0,0 +1,72 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import * as axis_util from '../../../ops/axis_util';
+import {computeOptimalWindowSize} from '../../../ops/reduce_util';
+import {sumOutType} from '../../../types';
+import {sizeFromShape} from '../../../util';
+import {MathBackendWebGL} from '../backend_webgl';
+import {ReduceProgram} from '../reduce_gpu';
+import {reshape} from '../reshape';
+
+interface SumInputs extends NamedTensorInfoMap {
+  x: TensorInfo;
+}
+
+interface SumAttrs extends NamedAttrMap {
+  axes: number[];
+}
+
+export const sumImpl =
+    (x: TensorInfo, reduceShape: number[],
+     backend: MathBackendWebGL): TensorInfo => {
+      const inSize = sizeFromShape(reduceShape);
+      const xSize = sizeFromShape(x.shape);
+      const batchSize = xSize / inSize;
+
+      x = reshape(x, [batchSize, inSize], backend);
+
+      const windowSize = computeOptimalWindowSize(inSize);
+      const reduceInfo = {windowSize, inSize, batchSize};
+      const program = new ReduceProgram(reduceInfo, 'sum');
+      const output = backend.runWebGLProgram(program, [x], sumOutType(x.dtype));
+
+      if (output.shape[1] === 1) {
+        return output;
+      }
+      return sumImpl(output, reduceShape, backend);
+    };
+
+registerKernel({
+  kernelName: 'Sum',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, attrs, backend}) => {
+    const {x} = inputs as SumInputs;
+    const {axes} = attrs as SumAttrs;
+    const webglBackend = backend as MathBackendWebGL;
+
+    axis_util.assertAxesAreInnerMostDims('sum', axes, x.shape.length);
+
+    const [outShape, reduceShape] =
+        axis_util.computeOutAndReduceShapes(x.shape, axes);
+
+    const out = sumImpl(x, reduceShape, webglBackend);
+
+    return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
+  }
+});

From 209a6021b7fbcb2140efbed5dc09e56fa59826a8 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 10:23:11 -0500
Subject: [PATCH 59/84] getting correct answer

---
 tfjs-core/src/backends/webgl/all_kernels.ts   |  1 +
 tfjs-core/src/backends/webgl/kernels/Div.ts   | 52 +++++++++++++++++++
 .../src/backends/webgl/kernels/Softmax.ts     | 11 ++--
 3 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 tfjs-core/src/backends/webgl/kernels/Div.ts

diff --git a/tfjs-core/src/backends/webgl/all_kernels.ts b/tfjs-core/src/backends/webgl/all_kernels.ts
index 7ed744f8b5f..d24280981f1 100644
--- a/tfjs-core/src/backends/webgl/all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/all_kernels.ts
@@ -21,6 +21,7 @@
 import './square';
 import './kernels/Softmax';
 import './kernels/Max';
+import './kernels/Div';
 import './kernels/Sum';
 import './kernels/Sub';
 import './kernels/Exp';
diff --git a/tfjs-core/src/backends/webgl/kernels/Div.ts b/tfjs-core/src/backends/webgl/kernels/Div.ts
new file mode 100644
index 00000000000..c97a3a311f9
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/kernels/Div.ts
@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {env} from '../../../environment';
+import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {MathBackendWebGL} from '../backend_webgl';
+import * as binaryop_gpu from '../binaryop_gpu';
+import {BinaryOpProgram} from '../binaryop_gpu';
+import * as binaryop_packed_gpu from '../binaryop_packed_gpu';
+import {BinaryOpPackedProgram} from '../binaryop_packed_gpu';
+
+interface DivInputs extends NamedTensorInfoMap {
+  a: TensorInfo;
+  b: TensorInfo;
+}
+
+export const divImpl =
+    (a: TensorInfo, b: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
+      let program = new BinaryOpProgram(binaryop_gpu.DIV, a.shape, b.shape);
+      if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
+        program = new BinaryOpPackedProgram(
+            binaryop_packed_gpu.DIV, a.shape, b.shape);
+      }
+      const output = backend.runWebGLProgram(program, [a, b], 'float32');
+      return output;
+    };
+
+registerKernel({
+  kernelName: 'Div',
+  backendName: 'webgl',
+  kernelFunc: ({inputs, backend}) => {
+    const {a, b} = inputs as DivInputs;
+    const webglBackend = backend as MathBackendWebGL;
+    const out = divImpl(a, b, webglBackend);
+
+    return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
+  }
+});
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index e891f0b6d67..0a94d7f87ae 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -21,6 +21,7 @@ import * as axis_util from '../../../ops/axis_util';
 import {parseAxisParam} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
 
+import {divImpl} from './Div';
 import {expImpl} from './Exp';
 import {maxImpl} from './Max';
 import {subImpl} from './Sub';
@@ -44,22 +45,22 @@ registerKernel({
 
     const axes = parseAxisParam([dim], logits.shape);
 
-    const [outShape, reduceShape] =
+    const [, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
 
     const max = maxImpl(logits, reduceShape, webglBackend);
-    console.log(outShape);
 
     const subtracted = subImpl(logits, max, webglBackend);
     const exponentiated = expImpl(subtracted, webglBackend);
 
-    const out = sumImpl(exponentiated, reduceShape, webglBackend);
+    const summed = sumImpl(exponentiated, reduceShape, webglBackend);
+
+    const out = divImpl(exponentiated, summed, webglBackend);
 
     webglBackend.disposeData(max.dataId);
     webglBackend.disposeData(subtracted.dataId);
     webglBackend.disposeData(exponentiated.dataId);
-
-    console.log('RAN SUB');
+    webglBackend.disposeData(summed.dataId);
 
     return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
   }

From 69a9d07efbd543899e981cb3e1c41f2e9664a935 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 10:35:23 -0500
Subject: [PATCH 60/84] fix

---
 tfjs-core/src/backends/webgl/kernels/Div.ts | 2 +-
 tfjs-core/src/backends/webgl/kernels/Sum.ts | 8 ++++----
 tfjs-core/src/ops/softmax_test.ts           | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Div.ts b/tfjs-core/src/backends/webgl/kernels/Div.ts
index c97a3a311f9..2be55ee8966 100644
--- a/tfjs-core/src/backends/webgl/kernels/Div.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Div.ts
@@ -33,7 +33,7 @@ export const divImpl =
       let program = new BinaryOpProgram(binaryop_gpu.DIV, a.shape, b.shape);
       if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
         program = new BinaryOpPackedProgram(
-            binaryop_packed_gpu.DIV, a.shape, b.shape);
+            binaryop_packed_gpu.DIV, a.shape, b.shape, true);
       }
       const output = backend.runWebGLProgram(program, [a, b], 'float32');
       return output;
diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index 65609649ac7..96288497d59 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -33,7 +33,7 @@ interface SumAttrs extends NamedAttrMap {
 }
 
 export const sumImpl =
-    (x: TensorInfo, reduceShape: number[],
+    (x: TensorInfo, reduceShape: number[], outShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
       const inSize = sizeFromShape(reduceShape);
       const xSize = sizeFromShape(x.shape);
@@ -47,9 +47,9 @@ export const sumImpl =
       const output = backend.runWebGLProgram(program, [x], sumOutType(x.dtype));
 
       if (output.shape[1] === 1) {
-        return output;
+        return reshape(output, outShape, backend);
       }
-      return sumImpl(output, reduceShape, backend);
+      return sumImpl(output, reduceShape, outShape, backend);
     };
 
 registerKernel({
@@ -65,7 +65,7 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(x.shape, axes);
 
-    const out = sumImpl(x, reduceShape, webglBackend);
+    const out = sumImpl(x, reduceShape, outShape, webglBackend);
 
     return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
   }
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 43260409c4b..c1a7d8e834c 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,7 +22,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  fit('regular test', async () => {
+  it('regular test', async () => {
     console.log('TESTINGGGG');
     // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From 018392bbbba3025d04d7d23c8e93e3939e3b16b4 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 10:37:44 -0500
Subject: [PATCH 61/84] softmax

---
 tfjs-core/src/backends/webgl/kernels/Softmax.ts | 4 ++--
 tfjs-core/src/ops/softmax_test.ts               | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index 0a94d7f87ae..5263374e2a5 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -45,7 +45,7 @@ registerKernel({
 
     const axes = parseAxisParam([dim], logits.shape);
 
-    const [, reduceShape] =
+    const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
 
     const max = maxImpl(logits, reduceShape, webglBackend);
@@ -53,7 +53,7 @@ registerKernel({
     const subtracted = subImpl(logits, max, webglBackend);
     const exponentiated = expImpl(subtracted, webglBackend);
 
-    const summed = sumImpl(exponentiated, reduceShape, webglBackend);
+    const summed = sumImpl(exponentiated, reduceShape, outShape, webglBackend);
 
     const out = divImpl(exponentiated, summed, webglBackend);
 
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index c1a7d8e834c..43260409c4b 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,7 +22,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  it('regular test', async () => {
+  fit('regular test', async () => {
     console.log('TESTINGGGG');
     // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From 66aad43213f5d8b099320b50d0a17215eb86f27f Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 11:12:32 -0500
Subject: [PATCH 62/84] fix sum

---
 tfjs-core/src/backends/webgl/kernels/Sum.ts | 25 +++++++----
 tfjs-core/src/backends/webgl/reshape.ts     | 49 +++++++++------------
 tfjs-core/src/ops/softmax_test.ts           |  2 +-
 3 files changed, 39 insertions(+), 37 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index 96288497d59..d2b4de3d476 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -32,6 +32,21 @@ interface SumAttrs extends NamedAttrMap {
   axes: number[];
 }
 
+const reduce =
+    (x: TensorInfo, reduceShape: number[],
+     backend: MathBackendWebGL): TensorInfo => {
+      const [batchSize, inSize] = x.shape;
+      const windowSize = computeOptimalWindowSize(inSize);
+      const reduceInfo = {windowSize, inSize, batchSize};
+      const program = new ReduceProgram(reduceInfo, 'sum');
+      const output = backend.runWebGLProgram(program, [x], sumOutType(x.dtype));
+
+      if (output.shape[1] === 1) {
+        return output;
+      }
+      return reduce(output, reduceShape, backend);
+    };
+
 export const sumImpl =
     (x: TensorInfo, reduceShape: number[], outShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
@@ -41,15 +56,7 @@ export const sumImpl =
 
       x = reshape(x, [batchSize, inSize], backend);
 
-      const windowSize = computeOptimalWindowSize(inSize);
-      const reduceInfo = {windowSize, inSize, batchSize};
-      const program = new ReduceProgram(reduceInfo, 'sum');
-      const output = backend.runWebGLProgram(program, [x], sumOutType(x.dtype));
-
-      if (output.shape[1] === 1) {
-        return reshape(output, outShape, backend);
-      }
-      return sumImpl(output, reduceShape, outShape, backend);
+      return reshape(reduce(x, reduceShape, backend), outShape, backend);
     };
 
 registerKernel({
diff --git a/tfjs-core/src/backends/webgl/reshape.ts b/tfjs-core/src/backends/webgl/reshape.ts
index 00d187b092c..0c773823e49 100644
--- a/tfjs-core/src/backends/webgl/reshape.ts
+++ b/tfjs-core/src/backends/webgl/reshape.ts
@@ -22,34 +22,29 @@ import {MathBackendWebGL} from './backend_webgl';
 import {ReshapePackedProgram} from './reshape_packed_gpu';
 
 const packedReshape =
-    (input: TensorInfo, afterShape: number[], backend: MathBackendWebGL):
-        TensorInfo => {
-          const input3DShape = [
-            webgl_util.getBatchDim(input.shape),
-            ...webgl_util.getRowsCols(input.shape)
-          ] as [number, number, number];
-          const input3D: TensorInfo = {
-            dtype: input.dtype,
-            shape: input3DShape,
-            dataId: input.dataId
-          };
-          const afterShapeAs3D = [
-            webgl_util.getBatchDim(afterShape),
-            ...webgl_util.getRowsCols(afterShape)
-          ] as [number, number, number];
+    (input: TensorInfo, afterShape: number[],
+     backend: MathBackendWebGL): TensorInfo => {
+      const input3DShape = [
+        webgl_util.getBatchDim(input.shape),
+        ...webgl_util.getRowsCols(input.shape)
+      ] as [number, number, number];
+      const input3D: TensorInfo = {
+        dtype: input.dtype,
+        shape: input3DShape,
+        dataId: input.dataId
+      };
+      const afterShapeAs3D = [
+        webgl_util.getBatchDim(afterShape),
+        ...webgl_util.getRowsCols(afterShape)
+      ] as [number, number, number];
 
-          const program =
-              new ReshapePackedProgram(afterShapeAs3D, input3DShape);
-          const preventEagerUnpackingOfOutput = true;
-          const output = backend.runWebGLProgram(
-              program, [input3D], input.dtype, null /* customSetup */,
-              preventEagerUnpackingOfOutput);
-          return {
-            dataId: output.dataId,
-            shape: afterShape,
-            dtype: output.dtype
-          };
-        }
+      const program = new ReshapePackedProgram(afterShapeAs3D, input3DShape);
+      const preventEagerUnpackingOfOutput = true;
+      const output = backend.runWebGLProgram(
+          program, [input3D], input.dtype, null /* customSetup */,
+          preventEagerUnpackingOfOutput);
+      return {dataId: output.dataId, shape: afterShape, dtype: output.dtype};
+    };
 
 export const reshape =
     (x: TensorInfo, afterShape: number[],
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 43260409c4b..c1a7d8e834c 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,7 +22,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  fit('regular test', async () => {
+  it('regular test', async () => {
     console.log('TESTINGGGG');
     // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From c61a3095c439d47344625a7b3ea5026c1e87bfc7 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 11:16:28 -0500
Subject: [PATCH 63/84] extract reduce

---
 tfjs-core/src/backends/webgl/kernels/Sum.ts | 23 +++----------
 tfjs-core/src/backends/webgl/reduce.ts      | 38 +++++++++++++++++++++
 2 files changed, 43 insertions(+), 18 deletions(-)
 create mode 100644 tfjs-core/src/backends/webgl/reduce.ts

diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index d2b4de3d476..2d06f3d0ee9 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -17,11 +17,11 @@
 
 import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
-import {computeOptimalWindowSize} from '../../../ops/reduce_util';
 import {sumOutType} from '../../../types';
+
 import {sizeFromShape} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
-import {ReduceProgram} from '../reduce_gpu';
+import {reduce} from '../reduce';
 import {reshape} from '../reshape';
 
 interface SumInputs extends NamedTensorInfoMap {
@@ -32,21 +32,6 @@ interface SumAttrs extends NamedAttrMap {
   axes: number[];
 }
 
-const reduce =
-    (x: TensorInfo, reduceShape: number[],
-     backend: MathBackendWebGL): TensorInfo => {
-      const [batchSize, inSize] = x.shape;
-      const windowSize = computeOptimalWindowSize(inSize);
-      const reduceInfo = {windowSize, inSize, batchSize};
-      const program = new ReduceProgram(reduceInfo, 'sum');
-      const output = backend.runWebGLProgram(program, [x], sumOutType(x.dtype));
-
-      if (output.shape[1] === 1) {
-        return output;
-      }
-      return reduce(output, reduceShape, backend);
-    };
-
 export const sumImpl =
     (x: TensorInfo, reduceShape: number[], outShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
@@ -56,7 +41,9 @@ export const sumImpl =
 
       x = reshape(x, [batchSize, inSize], backend);
 
-      return reshape(reduce(x, reduceShape, backend), outShape, backend);
+      return reshape(
+          reduce(x, reduceShape, sumOutType(x.dtype), backend), outShape,
+          backend);
     };
 
 registerKernel({
diff --git a/tfjs-core/src/backends/webgl/reduce.ts b/tfjs-core/src/backends/webgl/reduce.ts
new file mode 100644
index 00000000000..afd748c5444
--- /dev/null
+++ b/tfjs-core/src/backends/webgl/reduce.ts
@@ -0,0 +1,38 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {TensorInfo} from '../../kernel_registry';
+import {computeOptimalWindowSize} from '../../ops/reduce_util';
+import {DataType} from '../../types';
+
+import {MathBackendWebGL} from './backend_webgl';
+import {ReduceProgram} from './reduce_gpu';
+
+export const reduce =
+    (x: TensorInfo, reduceShape: number[], dtype: DataType,
+     backend: MathBackendWebGL): TensorInfo => {
+      const [batchSize, inSize] = x.shape;
+      const windowSize = computeOptimalWindowSize(inSize);
+      const reduceInfo = {windowSize, inSize, batchSize};
+      const program = new ReduceProgram(reduceInfo, 'sum');
+      const output = backend.runWebGLProgram(program, [x], dtype);
+
+      if (output.shape[1] === 1) {
+        return output;
+      }
+      return reduce(output, reduceShape, dtype, backend);
+    };

From 107c384fc6edbf412a1e0f9e80e82d522ce43f56 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 11:20:04 -0500
Subject: [PATCH 64/84] share reduce code

---
 tfjs-core/src/backends/webgl/kernels/Max.ts   | 30 +++++++------------
 .../src/backends/webgl/kernels/Softmax.ts     |  2 +-
 tfjs-core/src/ops/softmax_test.ts             |  2 +-
 3 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Max.ts b/tfjs-core/src/backends/webgl/kernels/Max.ts
index fb4bd83fc19..4f0e546b544 100644
--- a/tfjs-core/src/backends/webgl/kernels/Max.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Max.ts
@@ -17,10 +17,9 @@
 
 import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
-import {computeOptimalWindowSize} from '../../../ops/reduce_util';
 import {sizeFromShape} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
-import {ReduceProgram} from '../reduce_gpu';
+import {reduce} from '../reduce';
 import {reshape} from '../reshape';
 
 interface MaxInputs extends NamedTensorInfoMap {
@@ -32,24 +31,17 @@ interface MaxAttrs extends NamedAttrMap {
 }
 
 export const maxImpl =
-    (x: TensorInfo, reduceShape: number[], backend: MathBackendWebGL):
-        TensorInfo => {
-          const inSize = sizeFromShape(reduceShape);
-          const xSize = sizeFromShape(x.shape);
-          const batchSize = xSize / inSize;
+    (x: TensorInfo, reduceShape: number[], outShape: number[],
+     backend: MathBackendWebGL): TensorInfo => {
+      const inSize = sizeFromShape(reduceShape);
+      const xSize = sizeFromShape(x.shape);
+      const batchSize = xSize / inSize;
 
-          x = reshape(x, [batchSize, inSize], backend);
+      x = reshape(x, [batchSize, inSize], backend);
 
-          const windowSize = computeOptimalWindowSize(inSize);
-          const reduceInfo = {windowSize, inSize, batchSize};
-          const program = new ReduceProgram(reduceInfo, 'max');
-          const output = backend.runWebGLProgram(program, [x], x.dtype);
-
-          if (output.shape[1] === 1) {
-            return output;
-          }
-          return maxImpl(output, reduceShape, backend);
-        };
+      return reshape(
+          reduce(x, reduceShape, x.dtype, backend), outShape, backend);
+    };
 
 registerKernel({
   kernelName: 'Max',
@@ -64,7 +56,7 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(x.shape, axes);
 
-    const out = maxImpl(x, reduceShape, webglBackend);
+    const out = maxImpl(x, reduceShape, outShape, webglBackend);
 
     return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
   }
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index 5263374e2a5..e286a623de0 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -48,7 +48,7 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
 
-    const max = maxImpl(logits, reduceShape, webglBackend);
+    const max = maxImpl(logits, reduceShape, outShape, webglBackend);
 
     const subtracted = subImpl(logits, max, webglBackend);
     const exponentiated = expImpl(subtracted, webglBackend);
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index c1a7d8e834c..43260409c4b 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,7 +22,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  it('regular test', async () => {
+  fit('regular test', async () => {
     console.log('TESTINGGGG');
     // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From 0092d25ec4bb2b3a2ddcd222c2bceed0538cb206 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 13:16:08 -0500
Subject: [PATCH 65/84] setup

---
 tfjs-core/src/backends/webgl/backend_webgl.ts  | 10 ++++++++++
 tfjs-core/src/backends/webgl/kernels/Sum.ts    |  4 ++++
 tfjs-core/src/backends/webgl/reduce.ts         |  7 +++++++
 tfjs-core/src/backends/webgl/reshape.ts        |  2 ++
 tfjs-core/src/backends/webgl/webgl_ops_test.ts |  2 +-
 tfjs-core/src/engine.ts                        |  5 +++++
 tfjs-core/src/ops/softmax_test.ts              |  2 +-
 7 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index f367284a2a7..906314cc20c 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -577,19 +577,25 @@ export class MathBackendWebGL extends KernelBackend {
   private pendingDeletes = 0;
 
   disposeData(dataId: DataId): void {
+    console.log('ATTEMPTING TO RELEASE DATA');
     if (this.pendingDisposal.has(dataId)) {
+      console.log('pending disposal');
       return;
     }
     if (this.pendingRead.has(dataId)) {
+      console.log('pending read');
       this.pendingDisposal.add(dataId);
       this.pendingDeletes++;
       return;
     }
     // No-op if already disposed.
     if (!this.texData.has(dataId)) {
+      console.log('already disposed');
       return;
     }
 
+    console.log('RELEASE DATA');
+
     this.releaseGPUData(dataId);
     const {complexTensors} = this.texData.get(dataId);
     if (complexTensors != null) {
@@ -2496,6 +2502,7 @@ export class MathBackendWebGL extends KernelBackend {
       program: GPGPUProgram, inputs: TensorInfo[], outputDtype: DataType,
       customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => void,
       preventEagerUnpackingOfOutput = false): TensorInfo {
+    console.log('run webgl program', program.constructor.name);
     const output = this.makeTensorInfo(program.outputShape, outputDtype);
     const outData = this.texData.get(output.dataId);
     if (program.packedOutput) {
@@ -2579,10 +2586,12 @@ export class MathBackendWebGL extends KernelBackend {
         savedInput.shape = targetShape;
       }
 
+      console.log('maybe upload input...');
       this.uploadToGPU(input.dataId);
       return {shape: input.shape, texData, isUniform: false};
     });
 
+    console.log('maybe upload output....');
     this.uploadToGPU(output.dataId);
     const outputData:
         TensorData = {shape: output.shape, texData: outData, isUniform: false};
@@ -2704,6 +2713,7 @@ export class MathBackendWebGL extends KernelBackend {
       // Array is already on GPU. No-op.
       return;
     }
+    console.log('UPLOAD TO GPU');
     const shouldTimeProgram = this.activeTimers != null;
     let start: number;
     if (shouldTimeProgram) {
diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index 2d06f3d0ee9..0602676e7ce 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -35,6 +35,7 @@ interface SumAttrs extends NamedAttrMap {
 export const sumImpl =
     (x: TensorInfo, reduceShape: number[], outShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
+      console.log('SUM IMPL');
       const inSize = sizeFromShape(reduceShape);
       const xSize = sizeFromShape(x.shape);
       const batchSize = xSize / inSize;
@@ -50,6 +51,7 @@ registerKernel({
   kernelName: 'Sum',
   backendName: 'webgl',
   kernelFunc: ({inputs, attrs, backend}) => {
+    console.log('SUM KERNEL FUNC');
     const {x} = inputs as SumInputs;
     const {axes} = attrs as SumAttrs;
     const webglBackend = backend as MathBackendWebGL;
@@ -61,6 +63,8 @@ registerKernel({
 
     const out = sumImpl(x, reduceShape, outShape, webglBackend);
 
+    webglBackend.disposeData(x);
+
     return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
   }
 });
diff --git a/tfjs-core/src/backends/webgl/reduce.ts b/tfjs-core/src/backends/webgl/reduce.ts
index afd748c5444..c0a7799298a 100644
--- a/tfjs-core/src/backends/webgl/reduce.ts
+++ b/tfjs-core/src/backends/webgl/reduce.ts
@@ -25,14 +25,21 @@ import {ReduceProgram} from './reduce_gpu';
 export const reduce =
     (x: TensorInfo, reduceShape: number[], dtype: DataType,
      backend: MathBackendWebGL): TensorInfo => {
+      console.log('REDUCING -------------------------');
       const [batchSize, inSize] = x.shape;
       const windowSize = computeOptimalWindowSize(inSize);
       const reduceInfo = {windowSize, inSize, batchSize};
       const program = new ReduceProgram(reduceInfo, 'sum');
       const output = backend.runWebGLProgram(program, [x], dtype);
+      console.log('FINISHED COMPUTING REDUCE');
+      // I think the problem is that somehow, when recursing the input tensor is
+      // deemed to already be disposed, even though it's not...
+
+      backend.disposeData(x);
 
       if (output.shape[1] === 1) {
         return output;
       }
+
       return reduce(output, reduceShape, dtype, backend);
     };
diff --git a/tfjs-core/src/backends/webgl/reshape.ts b/tfjs-core/src/backends/webgl/reshape.ts
index 0c773823e49..2b3e47b3ef0 100644
--- a/tfjs-core/src/backends/webgl/reshape.ts
+++ b/tfjs-core/src/backends/webgl/reshape.ts
@@ -24,6 +24,7 @@ import {ReshapePackedProgram} from './reshape_packed_gpu';
 const packedReshape =
     (input: TensorInfo, afterShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
+      console.log('PACKED RESHAPE');
       const input3DShape = [
         webgl_util.getBatchDim(input.shape),
         ...webgl_util.getRowsCols(input.shape)
@@ -49,6 +50,7 @@ const packedReshape =
 export const reshape =
     (x: TensorInfo, afterShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
+      console.log('RESHAPE');
       const xTexData = backend.texData.get(x.dataId);
       if (xTexData.isPacked && !webgl_util.isReshapeFree(x.shape, afterShape) &&
           !(xTexData.texture !== null &&
diff --git a/tfjs-core/src/backends/webgl/webgl_ops_test.ts b/tfjs-core/src/backends/webgl/webgl_ops_test.ts
index 3dc5c1ebab1..e5db4bb177c 100644
--- a/tfjs-core/src/backends/webgl/webgl_ops_test.ts
+++ b/tfjs-core/src/backends/webgl/webgl_ops_test.ts
@@ -444,7 +444,7 @@ describeWithFlags('gramSchmidt-non-tiny', WEBGL_ENVS, () => {
 });
 
 describeWithFlags('matmul webgl-only', WEBGL_ENVS, () => {
-  it('Matrix times vector, large matrix', async () => {
+  fit('Matrix times vector, large matrix', async () => {
     const maxTexSize = 16000;
     const sharedDim = maxTexSize + 4;
     const matrix = tf.buffer<Rank.R2>([2, sharedDim], 'float32');
diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts
index 7f491b89bc3..ac2b6c40473 100644
--- a/tfjs-core/src/engine.ts
+++ b/tfjs-core/src/engine.ts
@@ -521,6 +521,8 @@ export class Engine implements TensorTracker, DataMover {
         this.state.numDataMovesStack[this.state.numDataMovesStack.length - 1];
     const dataIdsLeaked =
         numDataIdsAfter - numDataIdsBefore - numOutputDataIds - numMoves;
+
+    console.log(numDataIdsAfter, numDataIdsBefore, numOutputDataIds, numMoves);
     if (dataIdsLeaked > 0) {
       throw new Error(
           `Backend '${this.backendName}' has an internal memory leak ` +
@@ -566,10 +568,13 @@ export class Engine implements TensorTracker, DataMover {
     let out: TensorInfo|TensorInfo[];
     if (kernel != null) {
       kernelFunc = () => {
+        console.log('STARTING TO TRACK');
         const numDataIdsBefore = this.backend.numDataIds();
         out = kernel.kernelFunc({inputs, attrs, backend: this.backend});
         const outInfos = Array.isArray(out) ? out : [out];
         if (this.shouldCheckForMemLeaks()) {
+          console.log('CHECK KERNEL', numDataIdsBefore);
+          console.log(outInfos);
           this.checkKernelForMemLeak(kernelName, numDataIdsBefore, outInfos);
         }
         const outTensors = outInfos.map(
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 43260409c4b..c1a7d8e834c 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,7 +22,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  fit('regular test', async () => {
+  it('regular test', async () => {
     console.log('TESTINGGGG');
     // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From c42416dd49f00bf24143233c41802c908f52c744 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 14:36:31 -0500
Subject: [PATCH 66/84] dataid

---
 tfjs-core/src/backends/webgl/kernels/Sum.ts | 1 -
 tfjs-core/src/backends/webgl/reduce.ts      | 6 +-----
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index 0602676e7ce..a8641cf8f29 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -35,7 +35,6 @@ interface SumAttrs extends NamedAttrMap {
 export const sumImpl =
     (x: TensorInfo, reduceShape: number[], outShape: number[],
      backend: MathBackendWebGL): TensorInfo => {
-      console.log('SUM IMPL');
       const inSize = sizeFromShape(reduceShape);
       const xSize = sizeFromShape(x.shape);
       const batchSize = xSize / inSize;
diff --git a/tfjs-core/src/backends/webgl/reduce.ts b/tfjs-core/src/backends/webgl/reduce.ts
index c0a7799298a..4cee37841e2 100644
--- a/tfjs-core/src/backends/webgl/reduce.ts
+++ b/tfjs-core/src/backends/webgl/reduce.ts
@@ -25,17 +25,13 @@ import {ReduceProgram} from './reduce_gpu';
 export const reduce =
     (x: TensorInfo, reduceShape: number[], dtype: DataType,
      backend: MathBackendWebGL): TensorInfo => {
-      console.log('REDUCING -------------------------');
       const [batchSize, inSize] = x.shape;
       const windowSize = computeOptimalWindowSize(inSize);
       const reduceInfo = {windowSize, inSize, batchSize};
       const program = new ReduceProgram(reduceInfo, 'sum');
       const output = backend.runWebGLProgram(program, [x], dtype);
-      console.log('FINISHED COMPUTING REDUCE');
-      // I think the problem is that somehow, when recursing the input tensor is
-      // deemed to already be disposed, even though it's not...
 
-      backend.disposeData(x);
+      backend.disposeData(x.dataId);
 
       if (output.shape[1] === 1) {
         return output;

From bf4b9e71afd12d4cee434d7547037f37376a35c9 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 19 Feb 2020 14:37:58 -0500
Subject: [PATCH 67/84] remove logs

---
 tfjs-core/src/backends/webgl/backend_webgl.ts | 9 ---------
 tfjs-core/src/backends/webgl/kernels/Sum.ts   | 1 -
 2 files changed, 10 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 906314cc20c..93e6c57a4c4 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -577,25 +577,19 @@ export class MathBackendWebGL extends KernelBackend {
   private pendingDeletes = 0;
 
   disposeData(dataId: DataId): void {
-    console.log('ATTEMPTING TO RELEASE DATA');
     if (this.pendingDisposal.has(dataId)) {
-      console.log('pending disposal');
       return;
     }
     if (this.pendingRead.has(dataId)) {
-      console.log('pending read');
       this.pendingDisposal.add(dataId);
       this.pendingDeletes++;
       return;
     }
     // No-op if already disposed.
     if (!this.texData.has(dataId)) {
-      console.log('already disposed');
       return;
     }
 
-    console.log('RELEASE DATA');
-
     this.releaseGPUData(dataId);
     const {complexTensors} = this.texData.get(dataId);
     if (complexTensors != null) {
@@ -2502,7 +2496,6 @@ export class MathBackendWebGL extends KernelBackend {
       program: GPGPUProgram, inputs: TensorInfo[], outputDtype: DataType,
       customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => void,
       preventEagerUnpackingOfOutput = false): TensorInfo {
-    console.log('run webgl program', program.constructor.name);
     const output = this.makeTensorInfo(program.outputShape, outputDtype);
     const outData = this.texData.get(output.dataId);
     if (program.packedOutput) {
@@ -2586,12 +2579,10 @@ export class MathBackendWebGL extends KernelBackend {
         savedInput.shape = targetShape;
       }
 
-      console.log('maybe upload input...');
       this.uploadToGPU(input.dataId);
       return {shape: input.shape, texData, isUniform: false};
     });
 
-    console.log('maybe upload output....');
     this.uploadToGPU(output.dataId);
     const outputData:
         TensorData = {shape: output.shape, texData: outData, isUniform: false};
diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index a8641cf8f29..7f51ec5f1a6 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -50,7 +50,6 @@ registerKernel({
   kernelName: 'Sum',
   backendName: 'webgl',
   kernelFunc: ({inputs, attrs, backend}) => {
-    console.log('SUM KERNEL FUNC');
     const {x} = inputs as SumInputs;
     const {axes} = attrs as SumAttrs;
     const webglBackend = backend as MathBackendWebGL;

From 1e9da7c51fb2b55426c404c9c567ac60e2f3586f Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 18 Mar 2020 12:06:58 -0400
Subject: [PATCH 68/84] merge conflict

---
 tfjs-core/src/backends/cpu/kernels/Square.ts | 40 ++++++++++++++++++++
 tfjs-core/src/ops/softmax_test.ts            |  2 +-
 tfjs-core/src/util.ts                        | 32 ----------------
 3 files changed, 41 insertions(+), 33 deletions(-)
 create mode 100644 tfjs-core/src/backends/cpu/kernels/Square.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/Square.ts b/tfjs-core/src/backends/cpu/kernels/Square.ts
new file mode 100644
index 00000000000..e1cf21ed320
--- /dev/null
+++ b/tfjs-core/src/backends/cpu/kernels/Square.ts
@@ -0,0 +1,40 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. 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 {Square, SquareInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+export const squareConfig: KernelConfig = {
+  kernelName: Square,
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend}) => {
+    const {x} = inputs as SquareInputs;
+    const cpuBackend = backend as MathBackendCPU;
+    assertNotComplex(x, 'square');
+
+    const values = cpuBackend.data.get(x.dataId).values as Float32Array;
+    const newValues = new Float32Array(values.length);
+    for (let i = 0; i < values.length; ++i) {
+      const value = values[i];
+      newValues[i] = value * value;
+    }
+    const dataId = cpuBackend.write(newValues, x.shape, x.dtype);
+    return {dataId, shape: x.shape, dtype: x.dtype};
+  }
+};
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index c1a7d8e834c..43260409c4b 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,7 +22,7 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  it('regular test', async () => {
+  fit('regular test', async () => {
     console.log('TESTINGGGG');
     // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
diff --git a/tfjs-core/src/util.ts b/tfjs-core/src/util.ts
index dcf9bb8c7fb..438445dd452 100644
--- a/tfjs-core/src/util.ts
+++ b/tfjs-core/src/util.ts
@@ -548,38 +548,6 @@ export function computeStrides(shape: number[]): number[] {
   return strides;
 }
 
-export function locToIndex(locs: number[], strides: number[]): number {
-  const rank = locs.length;
-  if (rank === 0) {
-    return 0;
-  }
-  if (rank === 1) {
-    return locs[0];
-  }
-  let index = locs[rank - 1];
-  for (let i = 0; i < rank - 1; ++i) {
-    index += strides[i] * locs[i];
-  }
-  return index;
-}
-
-export function indexToLoc(index: number, strides: number[]): number[] {
-  const rank = strides.length + 1;
-  if (rank === 0) {
-    return [];
-  }
-  if (rank === 1) {
-    return [index];
-  }
-  const locs = new Array(rank);
-  for (let i = 0; i < locs.length - 1; ++i) {
-    locs[i] = Math.floor(index / strides[i]);
-    index -= locs[i] * strides[i];
-  }
-  locs[locs.length - 1] = index;
-  return locs;
-}
-
 export function toTypedArray(
     a: TensorLike, dtype: DataType, debugMode: boolean): TypedArray {
   if (dtype === 'string') {

From d07b6458e4f627681f53ae5b4f62cc2b40c4fca2 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 18 Mar 2020 12:16:53 -0400
Subject: [PATCH 69/84] build

---
 tfjs-core/src/backends/cpu/kernels/binary_impl.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
index d044bedee81..7d7a779de58 100644
--- a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
@@ -35,15 +35,15 @@ export const createBinaryOp = (op: any) =>
         }
       } else {
         for (let i = 0; i < outValues.length; ++i) {
-          const loc = util.indexToLoc(i, outStrides);
+          const loc = util.indexToLoc(i, outShape.length, outStrides);
 
           const aLoc = loc.slice(-aShape.length);
           aBroadcastDims.forEach(d => aLoc[d] = 0);
-          const aIndex = util.locToIndex(aLoc, aStrides);
+          const aIndex = util.locToIndex(aLoc, aShape.length, aStrides);
 
           const bLoc = loc.slice(-bShape.length);
           bBroadcastDims.forEach(d => bLoc[d] = 0);
-          const bIndex = util.locToIndex(bLoc, bStrides);
+          const bIndex = util.locToIndex(bLoc, bShape.length, bStrides);
 
           outValues[i] = op(a[aIndex], b[bIndex]);
         }

From 67e7153e85fee12470ba02e250de5edfd93e2026 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Wed, 18 Mar 2020 13:16:20 -0400
Subject: [PATCH 70/84] add cpu div

---
 tfjs-core/src/backends/cpu/backend_cpu.ts     | 12 ++---
 tfjs-core/src/backends/cpu/kernels/Div.ts     | 44 ++++++++++++++++++-
 .../src/backends/cpu/register_all_kernels.ts  |  5 +--
 tfjs-core/src/kernel_names.ts                 |  3 ++
 tfjs-core/src/ops/softmax_test.ts             | 21 ++++++---
 5 files changed, 67 insertions(+), 18 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/backend_cpu.ts b/tfjs-core/src/backends/cpu/backend_cpu.ts
index ee2a1e6772a..851f118fd55 100644
--- a/tfjs-core/src/backends/cpu/backend_cpu.ts
+++ b/tfjs-core/src/backends/cpu/backend_cpu.ts
@@ -482,13 +482,13 @@ export class MathBackendCPU extends KernelBackend {
         (aValue, bValue) => aValue * bValue);
   }
 
-  realDivide(a: Tensor, b: Tensor): Tensor {
-    assertNotComplex([a, b], 'realDivide');
+  // realDivide(a: Tensor, b: Tensor): Tensor {
+  //   assertNotComplex([a, b], 'realDivide');
 
-    const op = (a: number, b: number) => a / b;
-    const outputDtype = 'float32';
-    return this.broadcastedBinaryOp(a, b, outputDtype, op);
-  }
+  //   const op = (a: number, b: number) => a / b;
+  //   const outputDtype = 'float32';
+  //   return this.broadcastedBinaryOp(a, b, outputDtype, op);
+  // }
 
   floorDiv(a: Tensor, b: Tensor): Tensor {
     assertNotComplex([a, b], 'floorDiv');
diff --git a/tfjs-core/src/backends/cpu/kernels/Div.ts b/tfjs-core/src/backends/cpu/kernels/Div.ts
index 3efc0ed7453..528473a9faa 100644
--- a/tfjs-core/src/backends/cpu/kernels/Div.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Div.ts
@@ -15,7 +15,47 @@
  * =============================================================================
  */
 
-import {registerBinaryKernel} from './binary_kernel';
+// import {registerBinaryKernel} from './binary_kernel';
+// registerBinaryKernel('Div', div);
+
+import {Div, DivInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
+import * as broadcast_util from '../../../ops/broadcast_util';
+import {TypedArray} from '../../../types';
+import {sizeFromShape} from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+// import {broadcastedBinaryOp} from '../utils/kernel_utils';
+
 import {div} from './div_impl';
 
-registerBinaryKernel('Div', div);
+export const divConfig: KernelConfig = {
+  kernelName: Div,
+  backendName: 'cpu',
+  kernelFunc: ({inputs, backend}) => {
+    const {a, b} = inputs as DivInputs;
+    const cpuBackend = backend as MathBackendCPU;
+    assertNotComplex([a, b], Div);
+
+    const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
+    const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;
+
+    // const [resultData, resultShape] = broadcastedBinaryOp(
+    //     a.shape, b.shape, aVals, bVals, a.dtype, (aVal, bVal) => {
+    //       const diff = aVal - bVal;
+    //       return diff * diff;
+    //     });
+
+    const outShape =
+        broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
+    const outValues = new Float32Array(sizeFromShape(outShape));
+
+    const result = div(aVals, a.shape, bVals, b.shape, outValues, outShape);
+
+    const dataId = cpuBackend.write(result, outShape, a.dtype);
+    return {dataId, shape: outShape, dtype: a.dtype};
+
+    // const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
+    // return {dataId, shape: resultShape, dtype: a.dtype};
+  }
+};
diff --git a/tfjs-core/src/backends/cpu/register_all_kernels.ts b/tfjs-core/src/backends/cpu/register_all_kernels.ts
index 9516f9af311..4d7f2e983cf 100644
--- a/tfjs-core/src/backends/cpu/register_all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/register_all_kernels.ts
@@ -19,15 +19,14 @@
 // the contents of this file and import only the kernels that are needed.
 import {KernelConfig, registerKernel} from '../../kernel_registry';
 
+import {divConfig} from './kernels/Div';
 import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5';
 import {squareConfig} from './kernels/Square';
 import {squaredDifferenceConfig} from './kernels/SquaredDifference';
 
 // List all kernel configs here
 const kernelConfigs: KernelConfig[] = [
-  nonMaxSuppressionV5Config,
-  squareConfig,
-  squaredDifferenceConfig,
+  nonMaxSuppressionV5Config, squareConfig, squaredDifferenceConfig, divConfig
 ];
 
 for (const kernelConfig of kernelConfigs) {
diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts
index 2f027e57b1f..c6e764a45a2 100644
--- a/tfjs-core/src/kernel_names.ts
+++ b/tfjs-core/src/kernel_names.ts
@@ -21,6 +21,9 @@
 import {NamedTensorInfoMap} from './kernel_registry';
 import {PixelData} from './types';
 
+export const Div = 'Div';
+export type DivInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
+
 export const SquaredDifference = 'SquaredDifference';
 export type SquaredDifferenceInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
 
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 43260409c4b..2bfabf8bb23 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,21 +15,28 @@
  * =============================================================================
  */
 
-// import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
-import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
+import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
+// import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', WEBGL_ENVS, () => {
+describeWithFlags('softmax', CPU_ENVS, () => {
   fit('regular test', async () => {
     console.log('TESTINGGGG');
-    // const y = tf.softmax(tf.tensor1d([2, 1, 3]));
+    const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
-    // expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
+    expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
     // expectArraysClose(await y.sum().data(), 1);
-    const x = tf.tensor1d([1, 2, 3]);
-    x.softmax().print();
+    // const x = tf.tensor1d([1, 2, 3]);
+    // x.softmax().print();
+  });
+
+  fit('div', async () => {
+    console.log('TEST DIV');
+    const c = tf.div(tf.tensor1d([2, 1, 3]), tf.tensor1d([0.5, 0.5, 0.5]));
+    const data = await c.data();
+    console.log(data);
   });
 });
 

From a505261e9233d8834f93e3fc033bad336812c423 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 19 Mar 2020 09:46:45 -0400
Subject: [PATCH 71/84] clean

---
 tfjs-core/src/backends/cpu/kernels/Div.ts     | 26 ++--------
 .../src/backends/cpu/kernels/binary_impl.ts   | 49 ++++++++++++-------
 tfjs-core/src/backends/cpu/kernels/softmax.ts | 27 +++-------
 3 files changed, 42 insertions(+), 60 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/kernels/Div.ts b/tfjs-core/src/backends/cpu/kernels/Div.ts
index 528473a9faa..02ee847c11b 100644
--- a/tfjs-core/src/backends/cpu/kernels/Div.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Div.ts
@@ -15,17 +15,11 @@
  * =============================================================================
  */
 
-// import {registerBinaryKernel} from './binary_kernel';
-// registerBinaryKernel('Div', div);
-
 import {Div, DivInputs} from '../../../kernel_names';
 import {KernelConfig} from '../../../kernel_registry';
-import * as broadcast_util from '../../../ops/broadcast_util';
 import {TypedArray} from '../../../types';
-import {sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
-// import {broadcastedBinaryOp} from '../utils/kernel_utils';
 
 import {div} from './div_impl';
 
@@ -40,22 +34,10 @@ export const divConfig: KernelConfig = {
     const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
     const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;
 
-    // const [resultData, resultShape] = broadcastedBinaryOp(
-    //     a.shape, b.shape, aVals, bVals, a.dtype, (aVal, bVal) => {
-    //       const diff = aVal - bVal;
-    //       return diff * diff;
-    //     });
-
-    const outShape =
-        broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
-    const outValues = new Float32Array(sizeFromShape(outShape));
-
-    const result = div(aVals, a.shape, bVals, b.shape, outValues, outShape);
-
-    const dataId = cpuBackend.write(result, outShape, a.dtype);
-    return {dataId, shape: outShape, dtype: a.dtype};
+    const [resultData, resultShape] =
+        div(a.shape, b.shape, aVals, bVals, a.dtype);
 
-    // const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
-    // return {dataId, shape: resultShape, dtype: a.dtype};
+    const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
+    return {dataId, shape: resultShape, dtype: a.dtype};
   }
 };
diff --git a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
index 7d7a779de58..c1d37cc0ecd 100644
--- a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
@@ -15,38 +15,51 @@
  * =============================================================================
  */
 
-import * as broadcast_util from '../../../ops/broadcast_util';
-import {TypedArray} from '../../../types';
+import * as backend_util from '../../../backends/backend_util';
+
+import {DataType, NumericDataType, TypedArray} from '../../../types';
 import * as util from '../../../util';
 
-export const createBinaryOp = (op: any) =>
-    (a: TypedArray, aShape: number[], b: TypedArray, bShape: number[],
-     outValues: TypedArray, outShape: number[]): TypedArray => {
-      const aBroadcastDims = broadcast_util.getBroadcastDims(aShape, outShape);
-      const bBroadcastDims = broadcast_util.getBroadcastDims(bShape, outShape);
+export const createBinaryOp = (op: (a: number, b: number) => number) =>
+    (aShape: number[], bShape: number[], aVals: TypedArray, bVals: TypedArray,
+     dtype: DataType): [TypedArray, number[]] => {
+      const newShape = backend_util.assertAndGetBroadcastShape(aShape, bShape);
+
+      const resultRank = newShape.length;
+      const resultStrides = util.computeStrides(newShape);
+      const resultSize = util.sizeFromShape(newShape);
+
+      const result =
+          util.getTypedArrayFromDType(dtype as NumericDataType, resultSize);
+
+      const aRank = aShape.length;
+      const bRank = bShape.length;
 
       const aStrides = util.computeStrides(aShape);
       const bStrides = util.computeStrides(bShape);
-      const outStrides = util.computeStrides(outShape);
+
+      const aBroadcastDims = backend_util.getBroadcastDims(aShape, newShape);
+      const bBroadcastDims = backend_util.getBroadcastDims(bShape, newShape);
 
       if (aBroadcastDims.length + bBroadcastDims.length === 0) {
-        for (let i = 0; i < outValues.length; ++i) {
-          outValues[i] = op(a[i % a.length], b[i % b.length]);
+        for (let i = 0; i < result.length; ++i) {
+          result[i] = op(aVals[i % aVals.length], bVals[i % bVals.length]);
         }
       } else {
-        for (let i = 0; i < outValues.length; ++i) {
-          const loc = util.indexToLoc(i, outShape.length, outStrides);
+        for (let i = 0; i < result.length; ++i) {
+          const loc = util.indexToLoc(i, resultRank, resultStrides);
 
-          const aLoc = loc.slice(-aShape.length);
+          const aLoc = loc.slice(-aRank);
           aBroadcastDims.forEach(d => aLoc[d] = 0);
-          const aIndex = util.locToIndex(aLoc, aShape.length, aStrides);
+          const aIndex = util.locToIndex(aLoc, aRank, aStrides);
 
-          const bLoc = loc.slice(-bShape.length);
+          const bLoc = loc.slice(-bRank);
           bBroadcastDims.forEach(d => bLoc[d] = 0);
-          const bIndex = util.locToIndex(bLoc, bShape.length, bStrides);
+          const bIndex = util.locToIndex(bLoc, bRank, bStrides);
 
-          outValues[i] = op(a[aIndex], b[bIndex]);
+          result[i] = op(aVals[aIndex], bVals[bIndex]);
         }
       }
-      return outValues;
+
+      return [result, newShape];
     };
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index 9d330fa810d..c4b43d9b0c5 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -56,30 +56,17 @@ registerKernel({
 
     const expandedShape = axis_util.expandShapeToKeepDim(reduceOutShape, axes);
 
-    console.log('MAX LOGIT');
-    console.log(maxLogit);
+    const [aValues, aShape] =
+        sub(logits.shape, expandedShape, logitsValues, maxLogit, logits.dtype);
 
-    const a =
-        sub(logitsValues, logits.shape, maxLogit, expandedShape,
-            new Float32Array(sizeFromShape(logits.shape)), logits.shape);
-
-    console.log('subtract');
-    console.log(a);
-
-    const b = exp(a, new Float32Array(sizeFromShape(logits.shape)));
-    console.log('exp');
-    console.log(b);
+    const b = exp(aValues, new Float32Array(sizeFromShape(aShape)));
 
     const sumExp =
         sum(b, reduceShape, new Float32Array(sizeFromShape(reduceOutShape)));
-    console.log('sumexp');
-    console.log(sumExp);
-
-    const out =
-        div(b, logits.shape, sumExp, expandedShape,
-            new Float32Array(sizeFromShape(logits.shape)), logits.shape);
 
-    const dataId = cpuBackend.write(out, logits.shape, logits.dtype);
-    return {dataId, shape: logits.shape, dtype: logits.dtype};
+    const [resultData, resultShape] =
+        div(logits.shape, reduceShape, b, sumExp, logits.dtype);
+    const dataId = cpuBackend.write(resultData, resultShape, logits.dtype);
+    return {dataId, shape: resultShape, dtype: logits.dtype};
   }
 });

From 9ca3b6313f0a37758c310a826f118de39518dcc5 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 19 Mar 2020 10:11:01 -0400
Subject: [PATCH 72/84] move things around

---
 .../backends/cpu/kernels/SquaredDifference.ts | 14 ++--
 .../src/backends/cpu/kernels/binary_impl.ts   | 65 ------------------
 .../src/backends/cpu/kernels/div_impl.ts      |  2 +-
 .../src/backends/cpu/kernels/sub_impl.ts      |  2 +-
 .../src/backends/cpu/utils/kernel_utils.ts    | 68 +++++++++----------
 5 files changed, 44 insertions(+), 107 deletions(-)
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/binary_impl.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts b/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts
index a57bf1a156a..a9e4e174d66 100644
--- a/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts
+++ b/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts
@@ -20,7 +20,12 @@ import {KernelConfig} from '../../../kernel_registry';
 import {TypedArray} from '../../../types';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
-import {broadcastedBinaryOp} from '../utils/kernel_utils';
+import {createBinaryOp} from '../utils/kernel_utils';
+
+const squaredDifferenceImpl = createBinaryOp((aVal, bVal) => {
+  const diff = aVal - bVal;
+  return diff * diff;
+});
 
 export const squaredDifferenceConfig: KernelConfig = {
   kernelName: SquaredDifference,
@@ -33,11 +38,8 @@ export const squaredDifferenceConfig: KernelConfig = {
     const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
     const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;
 
-    const [resultData, resultShape] = broadcastedBinaryOp(
-        a.shape, b.shape, aVals, bVals, a.dtype, (aVal, bVal) => {
-          const diff = aVal - bVal;
-          return diff * diff;
-        });
+    const [resultData, resultShape] =
+        squaredDifferenceImpl(a.shape, b.shape, aVals, bVals, a.dtype);
 
     const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
     return {dataId, shape: resultShape, dtype: a.dtype};
diff --git a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts b/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
deleted file mode 100644
index c1d37cc0ecd..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/binary_impl.ts
+++ /dev/null
@@ -1,65 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 backend_util from '../../../backends/backend_util';
-
-import {DataType, NumericDataType, TypedArray} from '../../../types';
-import * as util from '../../../util';
-
-export const createBinaryOp = (op: (a: number, b: number) => number) =>
-    (aShape: number[], bShape: number[], aVals: TypedArray, bVals: TypedArray,
-     dtype: DataType): [TypedArray, number[]] => {
-      const newShape = backend_util.assertAndGetBroadcastShape(aShape, bShape);
-
-      const resultRank = newShape.length;
-      const resultStrides = util.computeStrides(newShape);
-      const resultSize = util.sizeFromShape(newShape);
-
-      const result =
-          util.getTypedArrayFromDType(dtype as NumericDataType, resultSize);
-
-      const aRank = aShape.length;
-      const bRank = bShape.length;
-
-      const aStrides = util.computeStrides(aShape);
-      const bStrides = util.computeStrides(bShape);
-
-      const aBroadcastDims = backend_util.getBroadcastDims(aShape, newShape);
-      const bBroadcastDims = backend_util.getBroadcastDims(bShape, newShape);
-
-      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
-        for (let i = 0; i < result.length; ++i) {
-          result[i] = op(aVals[i % aVals.length], bVals[i % bVals.length]);
-        }
-      } else {
-        for (let i = 0; i < result.length; ++i) {
-          const loc = util.indexToLoc(i, resultRank, resultStrides);
-
-          const aLoc = loc.slice(-aRank);
-          aBroadcastDims.forEach(d => aLoc[d] = 0);
-          const aIndex = util.locToIndex(aLoc, aRank, aStrides);
-
-          const bLoc = loc.slice(-bRank);
-          bBroadcastDims.forEach(d => bLoc[d] = 0);
-          const bIndex = util.locToIndex(bLoc, bRank, bStrides);
-
-          result[i] = op(aVals[aIndex], bVals[bIndex]);
-        }
-      }
-
-      return [result, newShape];
-    };
diff --git a/tfjs-core/src/backends/cpu/kernels/div_impl.ts b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
index ed14151713f..dd5e3453054 100644
--- a/tfjs-core/src/backends/cpu/kernels/div_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
@@ -15,6 +15,6 @@
  * =============================================================================
  */
 
-import {createBinaryOp} from './binary_impl';
+import {createBinaryOp} from '../utils/kernel_utils';
 
 export const div = createBinaryOp((a: number, b: number) => a / b);
diff --git a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
index 749e59da7e7..78bfa4b40a0 100644
--- a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
@@ -15,6 +15,6 @@
  * =============================================================================
  */
 
-import {createBinaryOp} from './binary_impl';
+import {createBinaryOp} from '../utils/kernel_utils';
 
 export const sub = createBinaryOp((a: number, b: number) => a - b);
diff --git a/tfjs-core/src/backends/cpu/utils/kernel_utils.ts b/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
index ce21517dba4..c1d37cc0ecd 100644
--- a/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
+++ b/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
@@ -16,50 +16,50 @@
  */
 
 import * as backend_util from '../../../backends/backend_util';
+
 import {DataType, NumericDataType, TypedArray} from '../../../types';
 import * as util from '../../../util';
 
-export function broadcastedBinaryOp(
-    aShape: number[], bShape: number[], aVals: TypedArray, bVals: TypedArray,
-    dtype: DataType,
-    op: (a: number, b: number) => number): [TypedArray, number[]] {
-  const newShape = backend_util.assertAndGetBroadcastShape(aShape, bShape);
+export const createBinaryOp = (op: (a: number, b: number) => number) =>
+    (aShape: number[], bShape: number[], aVals: TypedArray, bVals: TypedArray,
+     dtype: DataType): [TypedArray, number[]] => {
+      const newShape = backend_util.assertAndGetBroadcastShape(aShape, bShape);
 
-  const resultRank = newShape.length;
-  const resultStrides = util.computeStrides(newShape);
-  const resultSize = util.sizeFromShape(newShape);
+      const resultRank = newShape.length;
+      const resultStrides = util.computeStrides(newShape);
+      const resultSize = util.sizeFromShape(newShape);
 
-  const result =
-      util.getTypedArrayFromDType(dtype as NumericDataType, resultSize);
+      const result =
+          util.getTypedArrayFromDType(dtype as NumericDataType, resultSize);
 
-  const aRank = aShape.length;
-  const bRank = bShape.length;
+      const aRank = aShape.length;
+      const bRank = bShape.length;
 
-  const aStrides = util.computeStrides(aShape);
-  const bStrides = util.computeStrides(bShape);
+      const aStrides = util.computeStrides(aShape);
+      const bStrides = util.computeStrides(bShape);
 
-  const aBroadcastDims = backend_util.getBroadcastDims(aShape, newShape);
-  const bBroadcastDims = backend_util.getBroadcastDims(bShape, newShape);
+      const aBroadcastDims = backend_util.getBroadcastDims(aShape, newShape);
+      const bBroadcastDims = backend_util.getBroadcastDims(bShape, newShape);
 
-  if (aBroadcastDims.length + bBroadcastDims.length === 0) {
-    for (let i = 0; i < result.length; ++i) {
-      result[i] = op(aVals[i % aVals.length], bVals[i % bVals.length]);
-    }
-  } else {
-    for (let i = 0; i < result.length; ++i) {
-      const loc = util.indexToLoc(i, resultRank, resultStrides);
+      if (aBroadcastDims.length + bBroadcastDims.length === 0) {
+        for (let i = 0; i < result.length; ++i) {
+          result[i] = op(aVals[i % aVals.length], bVals[i % bVals.length]);
+        }
+      } else {
+        for (let i = 0; i < result.length; ++i) {
+          const loc = util.indexToLoc(i, resultRank, resultStrides);
 
-      const aLoc = loc.slice(-aRank);
-      aBroadcastDims.forEach(d => aLoc[d] = 0);
-      const aIndex = util.locToIndex(aLoc, aRank, aStrides);
+          const aLoc = loc.slice(-aRank);
+          aBroadcastDims.forEach(d => aLoc[d] = 0);
+          const aIndex = util.locToIndex(aLoc, aRank, aStrides);
 
-      const bLoc = loc.slice(-bRank);
-      bBroadcastDims.forEach(d => bLoc[d] = 0);
-      const bIndex = util.locToIndex(bLoc, bRank, bStrides);
+          const bLoc = loc.slice(-bRank);
+          bBroadcastDims.forEach(d => bLoc[d] = 0);
+          const bIndex = util.locToIndex(bLoc, bRank, bStrides);
 
-      result[i] = op(aVals[aIndex], bVals[bIndex]);
-    }
-  }
+          result[i] = op(aVals[aIndex], bVals[bIndex]);
+        }
+      }
 
-  return [result, newShape];
-}
+      return [result, newShape];
+    };

From 69b529344158f8d86a47aca598fc984744cd77d6 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 19 Mar 2020 10:30:21 -0400
Subject: [PATCH 73/84] max

---
 .../backends/cpu/kernels/{max.ts => Max.ts}   | 23 +++++++------------
 .../src/backends/cpu/kernels/max_impl.ts      |  9 ++++----
 tfjs-core/src/backends/cpu/kernels/softmax.ts |  5 ++--
 tfjs-core/src/kernel_names.ts                 |  6 +++++
 4 files changed, 21 insertions(+), 22 deletions(-)
 rename tfjs-core/src/backends/cpu/kernels/{max.ts => Max.ts} (77%)

diff --git a/tfjs-core/src/backends/cpu/kernels/max.ts b/tfjs-core/src/backends/cpu/kernels/Max.ts
similarity index 77%
rename from tfjs-core/src/backends/cpu/kernels/max.ts
rename to tfjs-core/src/backends/cpu/kernels/Max.ts
index 2a687b19284..1ac73aa4f85 100644
--- a/tfjs-core/src/backends/cpu/kernels/max.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Max.ts
@@ -15,27 +15,21 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {Max, MaxAttrs, MaxInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
 import {sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
-import {max} from './max_impl';
-
-interface MaxInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
 
-interface MaxAttrs extends NamedAttrMap {
-  axes: number[];
-}
+import {max} from './max_impl';
 
-registerKernel({
-  kernelName: 'Max',
+export const maxConfig: KernelConfig = {
+  kernelName: Max,
   backendName: 'cpu',
   kernelFunc: ({inputs, attrs, backend}) => {
     const {x} = inputs as MaxInputs;
-    const {axes} = attrs as MaxAttrs;
+    const {axes} = attrs as {} as MaxAttrs;
     const cpuBackend = backend as MathBackendCPU;
 
     assertNotComplex(x, 'max');
@@ -46,10 +40,9 @@ registerKernel({
         axis_util.computeOutAndReduceShapes(x.shape, axes);
 
     const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
-    const outValues = new Float32Array(sizeFromShape(outShape));
-    const result = max(xVals, reduceShape, outValues);
+    const result = max(xVals, sizeFromShape(reduceShape), outShape, x.dtype);
 
     const dataId = cpuBackend.write(result, outShape, x.dtype);
     return {dataId, shape: outShape, dtype: x.dtype};
   }
-});
+};
diff --git a/tfjs-core/src/backends/cpu/kernels/max_impl.ts b/tfjs-core/src/backends/cpu/kernels/max_impl.ts
index 49eee5d4ce2..f03a7a7c8d8 100644
--- a/tfjs-core/src/backends/cpu/kernels/max_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/max_impl.ts
@@ -15,13 +15,14 @@
  * =============================================================================
  */
 
-import {TypedArray} from '../../../types';
-import {sizeFromShape} from '../../../util';
+import {DataType, NumericDataType, TypedArray} from '../../../types';
+import * as util from '../../../util';
 
 export const max =
-    (x: TypedArray, reduceShape: number[], outValues: TypedArray):
+    (x: TypedArray, reduceSize: number, outShape: number[], dtype: DataType):
         TypedArray => {
-          const reduceSize = sizeFromShape(reduceShape);
+          const outValues = util.getTypedArrayFromDType(
+              dtype as NumericDataType, util.sizeFromShape(outShape));
 
           for (let i = 0; i < x.length; ++i) {
             const offset = i * reduceSize;
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index c4b43d9b0c5..8ca0fe966c8 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -50,9 +50,8 @@ registerKernel({
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
     const logitsValues =
         cpuBackend.data.get(logits.dataId).values as Float32Array;
-    const maxLogit =
-        max(logitsValues, reduceShape,
-            new Float32Array(sizeFromShape(reduceOutShape)));
+    const maxLogit = max(
+        logitsValues, sizeFromShape(reduceShape), reduceOutShape, logits.dtype);
 
     const expandedShape = axis_util.expandShapeToKeepDim(reduceOutShape, axes);
 
diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts
index c6e764a45a2..e7533bfc5f4 100644
--- a/tfjs-core/src/kernel_names.ts
+++ b/tfjs-core/src/kernel_names.ts
@@ -24,6 +24,12 @@ import {PixelData} from './types';
 export const Div = 'Div';
 export type DivInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
 
+export const Max = 'Max';
+export type MaxInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface MaxAttrs {
+  axes: number[];
+}
+
 export const SquaredDifference = 'SquaredDifference';
 export type SquaredDifferenceInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
 

From 26e1db2bd6020b1789201ab3ff5f7fc557e5d680 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 19 Mar 2020 10:37:15 -0400
Subject: [PATCH 74/84] exp

---
 tfjs-core/src/backends/cpu/kernels/Exp.ts      | 16 ++++++----------
 tfjs-core/src/backends/cpu/kernels/exp_impl.ts |  5 ++++-
 tfjs-core/src/backends/cpu/kernels/softmax.ts  |  4 ++--
 tfjs-core/src/kernel_names.ts                  |  3 +++
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/kernels/Exp.ts b/tfjs-core/src/backends/cpu/kernels/Exp.ts
index e834aa6e1fd..e607dc21938 100644
--- a/tfjs-core/src/backends/cpu/kernels/Exp.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Exp.ts
@@ -15,18 +15,14 @@
  * =============================================================================
  */
 
-import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
-import {sizeFromShape} from '../../../util';
+import {Exp, ExpInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
 import {MathBackendCPU} from '../backend_cpu';
 
 import {exp} from './exp_impl';
 
-interface ExpInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-registerKernel({
-  kernelName: 'Exp',
+export const expConfig: KernelConfig = {
+  kernelName: Exp,
   backendName: 'cpu',
   kernelFunc: ({inputs, backend}) => {
     const {x} = inputs as ExpInputs;
@@ -34,9 +30,9 @@ registerKernel({
 
     const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
 
-    const result = exp(xVals, new Float32Array(sizeFromShape(x.shape)));
+    const result = exp(xVals);
 
     const dataId = cpuBackend.write(result, x.shape, x.dtype);
     return {dataId, shape: x.shape, dtype: x.dtype};
   }
-});
+};
diff --git a/tfjs-core/src/backends/cpu/kernels/exp_impl.ts b/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
index e8d8dd5b35e..111ad85455e 100644
--- a/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
@@ -16,8 +16,11 @@
  */
 
 import {TypedArray} from '../../../types';
+import * as util from '../../../util';
+
+export const exp = (x: TypedArray): TypedArray => {
+  const outValues = util.getTypedArrayFromDType('float32', x.length);
 
-export const exp = (x: TypedArray, outValues: TypedArray): TypedArray => {
   for (let i = 0; i < x.length; ++i) {
     outValues[i] = Math.exp(x[i]);
   }
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index 8ca0fe966c8..7acba47e7dc 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -55,10 +55,10 @@ registerKernel({
 
     const expandedShape = axis_util.expandShapeToKeepDim(reduceOutShape, axes);
 
-    const [aValues, aShape] =
+    const [aValues,] =
         sub(logits.shape, expandedShape, logitsValues, maxLogit, logits.dtype);
 
-    const b = exp(aValues, new Float32Array(sizeFromShape(aShape)));
+    const b = exp(aValues);
 
     const sumExp =
         sum(b, reduceShape, new Float32Array(sizeFromShape(reduceOutShape)));
diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts
index e7533bfc5f4..bb8216ca236 100644
--- a/tfjs-core/src/kernel_names.ts
+++ b/tfjs-core/src/kernel_names.ts
@@ -24,6 +24,9 @@ import {PixelData} from './types';
 export const Div = 'Div';
 export type DivInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
 
+export const Exp = 'Exp';
+export type ExpInputs = Pick<NamedTensorInfoMap, 'x'>;
+
 export const Max = 'Max';
 export type MaxInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface MaxAttrs {

From 5607f5c5f07c78b3a74feb85902c249c887dfc9d Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 19 Mar 2020 15:43:43 -0400
Subject: [PATCH 75/84] sub

---
 tfjs-core/src/backends/cpu/kernels/Div.ts     | 27 ++--------
 .../backends/cpu/kernels/SquaredDifference.ts | 27 ++--------
 tfjs-core/src/backends/cpu/kernels/Sub.ts     |  6 ++-
 .../src/backends/cpu/kernels/binary_kernel.ts | 51 -------------------
 .../src/backends/cpu/utils/kernel_utils.ts    | 33 +++++++++++-
 tfjs-core/src/kernel_names.ts                 |  4 +-
 6 files changed, 45 insertions(+), 103 deletions(-)
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/binary_kernel.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/Div.ts b/tfjs-core/src/backends/cpu/kernels/Div.ts
index 02ee847c11b..16d06a54666 100644
--- a/tfjs-core/src/backends/cpu/kernels/Div.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Div.ts
@@ -15,29 +15,8 @@
  * =============================================================================
  */
 
-import {Div, DivInputs} from '../../../kernel_names';
-import {KernelConfig} from '../../../kernel_registry';
-import {TypedArray} from '../../../types';
-import {MathBackendCPU} from '../backend_cpu';
-import {assertNotComplex} from '../cpu_util';
+import {Div} from '../../../kernel_names';
+import {createBinaryKernelConfig} from '../utils/kernel_utils';
 
 import {div} from './div_impl';
-
-export const divConfig: KernelConfig = {
-  kernelName: Div,
-  backendName: 'cpu',
-  kernelFunc: ({inputs, backend}) => {
-    const {a, b} = inputs as DivInputs;
-    const cpuBackend = backend as MathBackendCPU;
-    assertNotComplex([a, b], Div);
-
-    const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
-    const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;
-
-    const [resultData, resultShape] =
-        div(a.shape, b.shape, aVals, bVals, a.dtype);
-
-    const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
-    return {dataId, shape: resultShape, dtype: a.dtype};
-  }
-};
+export const divConfig = createBinaryKernelConfig(Div, div);
diff --git a/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts b/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts
index a9e4e174d66..d89a0e71181 100644
--- a/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts
+++ b/tfjs-core/src/backends/cpu/kernels/SquaredDifference.ts
@@ -15,33 +15,14 @@
  * =============================================================================
  */
 
-import {SquaredDifference, SquaredDifferenceInputs} from '../../../kernel_names';
-import {KernelConfig} from '../../../kernel_registry';
-import {TypedArray} from '../../../types';
-import {MathBackendCPU} from '../backend_cpu';
-import {assertNotComplex} from '../cpu_util';
+import {SquaredDifference} from '../../../kernel_names';
 import {createBinaryOp} from '../utils/kernel_utils';
+import {createBinaryKernelConfig} from '../utils/kernel_utils';
 
 const squaredDifferenceImpl = createBinaryOp((aVal, bVal) => {
   const diff = aVal - bVal;
   return diff * diff;
 });
 
-export const squaredDifferenceConfig: KernelConfig = {
-  kernelName: SquaredDifference,
-  backendName: 'cpu',
-  kernelFunc: ({inputs, backend}) => {
-    const {a, b} = inputs as SquaredDifferenceInputs;
-    const cpuBackend = backend as MathBackendCPU;
-    assertNotComplex([a, b], SquaredDifference);
-
-    const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
-    const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;
-
-    const [resultData, resultShape] =
-        squaredDifferenceImpl(a.shape, b.shape, aVals, bVals, a.dtype);
-
-    const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
-    return {dataId, shape: resultShape, dtype: a.dtype};
-  }
-};
+export const squaredDifferenceConfig =
+    createBinaryKernelConfig(SquaredDifference, squaredDifferenceImpl);
diff --git a/tfjs-core/src/backends/cpu/kernels/Sub.ts b/tfjs-core/src/backends/cpu/kernels/Sub.ts
index d82fab58a23..96ed5e074fb 100644
--- a/tfjs-core/src/backends/cpu/kernels/Sub.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Sub.ts
@@ -15,7 +15,9 @@
  * =============================================================================
  */
 
-import {registerBinaryKernel} from './binary_kernel';
+import {Sub} from '../../../kernel_names';
+import {createBinaryKernelConfig} from '../utils/kernel_utils';
+
 import {sub} from './sub_impl';
 
-registerBinaryKernel('Sub', sub);
+export const subConfig = createBinaryKernelConfig(Sub, sub);
diff --git a/tfjs-core/src/backends/cpu/kernels/binary_kernel.ts b/tfjs-core/src/backends/cpu/kernels/binary_kernel.ts
deleted file mode 100644
index b7e528e3c73..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/binary_kernel.ts
+++ /dev/null
@@ -1,51 +0,0 @@
-/**
- * @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 {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
-import * as broadcast_util from '../../../ops/broadcast_util';
-import {upcastType} from '../../../types';
-import {sizeFromShape} from '../../../util';
-import {MathBackendCPU} from '../backend_cpu';
-
-export function registerBinaryKernel(kernelName: string, op: any) {
-  registerKernel({
-    kernelName,
-    backendName: 'cpu',
-    kernelFunc: ({inputs, backend}) => {
-      const {a, b} = inputs as BinaryInputs;
-      const cpuBackend = backend as MathBackendCPU;
-
-      const dtype = upcastType(a.dtype, b.dtype);
-      const outShape =
-          broadcast_util.assertAndGetBroadcastShape(a.shape, b.shape);
-      const outValues = new Float32Array(sizeFromShape(outShape));
-
-      const aVals = cpuBackend.data.get(a.dataId).values as Float32Array;
-      const bVals = cpuBackend.data.get(b.dataId).values as Float32Array;
-
-      const result = op(aVals, a.shape, bVals, b.shape, outValues, outShape);
-
-      const dataId = cpuBackend.write(result, outShape, dtype);
-      return {dataId, shape: outShape, dtype};
-    }
-  });
-}
-
-interface BinaryInputs extends NamedTensorInfoMap {
-  a: TensorInfo;
-  b: TensorInfo;
-}
diff --git a/tfjs-core/src/backends/cpu/utils/kernel_utils.ts b/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
index c1d37cc0ecd..b1eafa8c9f2 100644
--- a/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
+++ b/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
@@ -16,9 +16,40 @@
  */
 
 import * as backend_util from '../../../backends/backend_util';
-
+import {KernelConfig, NamedTensorInfoMap, TensorInfo} from '../../../kernel_registry';
 import {DataType, NumericDataType, TypedArray} from '../../../types';
 import * as util from '../../../util';
+import {MathBackendCPU} from '../backend_cpu';
+import {assertNotComplex} from '../cpu_util';
+
+interface BinaryInputs extends NamedTensorInfoMap {
+  a: TensorInfo;
+  b: TensorInfo;
+}
+
+export const createBinaryKernelConfig =
+    (name: string,
+     op: (
+         aShape: number[], bShape: number[], aVals: TypedArray,
+         bVals: TypedArray, dtype: DataType) => [TypedArray, number[]]):
+        KernelConfig => ({
+          kernelName: name,
+          backendName: 'cpu',
+          kernelFunc: ({inputs, backend}) => {
+            const {a, b} = inputs as BinaryInputs;
+            const cpuBackend = backend as MathBackendCPU;
+            assertNotComplex([a, b], name);
+
+            const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
+            const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;
+
+            const [resultData, resultShape] =
+                op(a.shape, b.shape, aVals, bVals, a.dtype);
+
+            const dataId = cpuBackend.write(resultData, resultShape, a.dtype);
+            return {dataId, shape: resultShape, dtype: a.dtype};
+          }
+        });
 
 export const createBinaryOp = (op: (a: number, b: number) => number) =>
     (aShape: number[], bShape: number[], aVals: TypedArray, bVals: TypedArray,
diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts
index bb8216ca236..b592c24245f 100644
--- a/tfjs-core/src/kernel_names.ts
+++ b/tfjs-core/src/kernel_names.ts
@@ -22,7 +22,6 @@ import {NamedTensorInfoMap} from './kernel_registry';
 import {PixelData} from './types';
 
 export const Div = 'Div';
-export type DivInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
 
 export const Exp = 'Exp';
 export type ExpInputs = Pick<NamedTensorInfoMap, 'x'>;
@@ -34,11 +33,12 @@ export interface MaxAttrs {
 }
 
 export const SquaredDifference = 'SquaredDifference';
-export type SquaredDifferenceInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
 
 export const Square = 'Square';
 export type SquareInputs = Pick<NamedTensorInfoMap, 'x'>;
 
+export const Sub = 'Sub';
+
 export const NonMaxSuppressionV5 = 'NonMaxSuppressionV5';
 export type NonMaxSuppressionV5Inputs =
     Pick<NamedTensorInfoMap, 'boxes'|'scores'>;

From 96877940aa13e6c263c55e093557999589fe55f0 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 08:50:18 -0400
Subject: [PATCH 76/84] add sum

---
 tfjs-core/src/backends/cpu/kernels/Sum.ts     | 22 ++++++-------------
 tfjs-core/src/backends/cpu/kernels/softmax.ts |  2 +-
 .../src/backends/cpu/kernels/sum_impl.ts      |  9 ++++----
 .../src/backends/cpu/utils/kernel_utils.ts    |  8 ++-----
 .../webgl/kernels/SquaredDifference.ts        |  4 ++--
 tfjs-core/src/kernel_names.ts                 |  8 +++++++
 tfjs-core/src/ops/squared_difference.ts       |  4 ++--
 7 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/kernels/Sum.ts b/tfjs-core/src/backends/cpu/kernels/Sum.ts
index 2b309ec077f..43c4930e8f9 100644
--- a/tfjs-core/src/backends/cpu/kernels/Sum.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Sum.ts
@@ -15,7 +15,8 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {Sum, SumAttrs, SumInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
 import {upcastType} from '../../../types';
 import {sizeFromShape} from '../../../util';
@@ -24,20 +25,12 @@ import {assertNotComplex} from '../cpu_util';
 
 import {sum} from './sum_impl';
 
-interface SumInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface SumAttrs extends NamedAttrMap {
-  axes: number[];
-}
-
-registerKernel({
-  kernelName: 'Sum',
+export const sumConfig: KernelConfig = {
+  kernelName: Sum,
   backendName: 'cpu',
   kernelFunc: ({inputs, attrs, backend}) => {
     const {x} = inputs as SumInputs;
-    const {axes} = attrs as SumAttrs;
+    const {axes} = attrs as {} as SumAttrs;
     const cpuBackend = backend as MathBackendCPU;
 
     assertNotComplex(x, 'sum');
@@ -49,10 +42,9 @@ registerKernel({
     const resultDtype = upcastType(x.dtype, 'int32');
 
     const xVals = cpuBackend.data.get(x.dataId).values as Float32Array;
-    const result =
-        sum(xVals, reduceShape, new Float32Array(sizeFromShape(outShape)));
+    const result = sum(xVals, sizeFromShape(reduceShape), outShape, x.dtype);
 
     const dataId = cpuBackend.write(result, outShape, resultDtype);
     return {dataId, shape: outShape, dtype: resultDtype};
   }
-});
+};
diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/softmax.ts
index 7acba47e7dc..bd2de9abb08 100644
--- a/tfjs-core/src/backends/cpu/kernels/softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/softmax.ts
@@ -61,7 +61,7 @@ registerKernel({
     const b = exp(aValues);
 
     const sumExp =
-        sum(b, reduceShape, new Float32Array(sizeFromShape(reduceOutShape)));
+        sum(b, sizeFromShape(reduceShape), reduceOutShape, logits.dtype);
 
     const [resultData, resultShape] =
         div(logits.shape, reduceShape, b, sumExp, logits.dtype);
diff --git a/tfjs-core/src/backends/cpu/kernels/sum_impl.ts b/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
index e0ba88d7295..36bcdce8c48 100644
--- a/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
+++ b/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
@@ -15,13 +15,14 @@
  * =============================================================================
  */
 
-import {TypedArray} from '../../../types';
-import {sizeFromShape} from '../../../util';
+import {DataType, NumericDataType, TypedArray} from '../../../types';
+import * as util from '../../../util';
 
 export const sum =
-    (x: TypedArray, reduceShape: number[], outValues: TypedArray):
+    (x: TypedArray, reduceSize: number, outShape: number[], dtype: DataType):
         TypedArray => {
-          const reduceSize = sizeFromShape(reduceShape);
+          const outValues = util.getTypedArrayFromDType(
+              dtype as NumericDataType, util.sizeFromShape(outShape));
 
           for (let i = 0; i < x.length; ++i) {
             const offset = i * reduceSize;
diff --git a/tfjs-core/src/backends/cpu/utils/kernel_utils.ts b/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
index b1eafa8c9f2..f9fa95e9903 100644
--- a/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
+++ b/tfjs-core/src/backends/cpu/utils/kernel_utils.ts
@@ -16,17 +16,13 @@
  */
 
 import * as backend_util from '../../../backends/backend_util';
-import {KernelConfig, NamedTensorInfoMap, TensorInfo} from '../../../kernel_registry';
+import {BinaryInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
 import {DataType, NumericDataType, TypedArray} from '../../../types';
 import * as util from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
-interface BinaryInputs extends NamedTensorInfoMap {
-  a: TensorInfo;
-  b: TensorInfo;
-}
-
 export const createBinaryKernelConfig =
     (name: string,
      op: (
diff --git a/tfjs-core/src/backends/webgl/kernels/SquaredDifference.ts b/tfjs-core/src/backends/webgl/kernels/SquaredDifference.ts
index 41463b2c1cb..9ff08538a74 100644
--- a/tfjs-core/src/backends/webgl/kernels/SquaredDifference.ts
+++ b/tfjs-core/src/backends/webgl/kernels/SquaredDifference.ts
@@ -16,7 +16,7 @@
  */
 
 import {env} from '../../../environment';
-import {SquaredDifference, SquaredDifferenceInputs} from '../../../kernel_names';
+import {BinaryInputs, SquaredDifference} from '../../../kernel_names';
 import {KernelConfig} from '../../../kernel_registry';
 import {MathBackendWebGL} from '../backend_webgl';
 import {BinaryOpProgram} from '../binaryop_gpu';
@@ -26,7 +26,7 @@ export const squaredDifferenceConfig: KernelConfig = {
   kernelName: SquaredDifference,
   backendName: 'webgl',
   kernelFunc: ({inputs, backend}) => {
-    const {a, b} = inputs as SquaredDifferenceInputs;
+    const {a, b} = inputs as BinaryInputs;
     const SQUARED_DIFFERENCE = 'return (a - b) * (a - b);';
     const webGLBackend = backend as MathBackendWebGL;
 
diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts
index b592c24245f..eedcc24980e 100644
--- a/tfjs-core/src/kernel_names.ts
+++ b/tfjs-core/src/kernel_names.ts
@@ -21,6 +21,8 @@
 import {NamedTensorInfoMap} from './kernel_registry';
 import {PixelData} from './types';
 
+export type BinaryInputs = Pick<NamedTensorInfoMap, 'a'|'b'>;
+
 export const Div = 'Div';
 
 export const Exp = 'Exp';
@@ -39,6 +41,12 @@ export type SquareInputs = Pick<NamedTensorInfoMap, 'x'>;
 
 export const Sub = 'Sub';
 
+export const Sum = 'Sum';
+export type SumInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface SumAttrs {
+  axes: number[];
+}
+
 export const NonMaxSuppressionV5 = 'NonMaxSuppressionV5';
 export type NonMaxSuppressionV5Inputs =
     Pick<NamedTensorInfoMap, 'boxes'|'scores'>;
diff --git a/tfjs-core/src/ops/squared_difference.ts b/tfjs-core/src/ops/squared_difference.ts
index 0653191326c..9a5ecef3641 100644
--- a/tfjs-core/src/ops/squared_difference.ts
+++ b/tfjs-core/src/ops/squared_difference.ts
@@ -16,7 +16,7 @@
  */
 
 import {ENGINE, ForwardFunc} from '../engine';
-import {SquaredDifference, SquaredDifferenceInputs} from '../kernel_names';
+import {BinaryInputs, SquaredDifference} from '../kernel_names';
 import {Tensor} from '../tensor';
 import {NamedTensorMap} from '../tensor_types';
 import {makeTypesMatch} from '../tensor_util';
@@ -74,7 +74,7 @@ function squaredDifference_<T extends Tensor>(
     return res;
   };
 
-  const inputs: SquaredDifferenceInputs = {a: $a, b: $b};
+  const inputs: BinaryInputs = {a: $a, b: $b};
   const attrs = {};
 
   const inputsToSave = [$a, $b];

From e9b00bfdd02f5624916f6ad7fb3401db4fd91d3c Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 09:59:08 -0400
Subject: [PATCH 77/84] cap

---
 .../src/backends/cpu/kernels/{softmax.ts => Softmax.ts}     | 0
 tfjs-core/src/ops/softmax_test.ts                           | 6 +++---
 2 files changed, 3 insertions(+), 3 deletions(-)
 rename tfjs-core/src/backends/cpu/kernels/{softmax.ts => Softmax.ts} (100%)

diff --git a/tfjs-core/src/backends/cpu/kernels/softmax.ts b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
similarity index 100%
rename from tfjs-core/src/backends/cpu/kernels/softmax.ts
rename to tfjs-core/src/backends/cpu/kernels/Softmax.ts
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 2bfabf8bb23..e87aa5e8a2a 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,13 +15,13 @@
  * =============================================================================
  */
 
-import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
-// import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
+// import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
+import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', CPU_ENVS, () => {
+describeWithFlags('softmax', WEBGL_ENVS, () => {
   fit('regular test', async () => {
     console.log('TESTINGGGG');
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));

From 26dadc3afddaf4438ed7bdcb76daa5b43f4a77b0 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 10:02:01 -0400
Subject: [PATCH 78/84] simplify

---
 tfjs-core/src/backends/cpu/kernels/Div.ts     |  3 ++-
 tfjs-core/src/backends/cpu/kernels/Softmax.ts |  2 +-
 .../src/backends/cpu/kernels/div_impl.ts      | 20 -------------------
 tfjs-core/src/ops/softmax_test.ts             |  6 +++---
 4 files changed, 6 insertions(+), 25 deletions(-)
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/div_impl.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/Div.ts b/tfjs-core/src/backends/cpu/kernels/Div.ts
index 16d06a54666..5483f92f857 100644
--- a/tfjs-core/src/backends/cpu/kernels/Div.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Div.ts
@@ -17,6 +17,7 @@
 
 import {Div} from '../../../kernel_names';
 import {createBinaryKernelConfig} from '../utils/kernel_utils';
+import {createBinaryOp} from '../utils/kernel_utils';
 
-import {div} from './div_impl';
+export const div = createBinaryOp((a: number, b: number) => a / b);
 export const divConfig = createBinaryKernelConfig(Div, div);
diff --git a/tfjs-core/src/backends/cpu/kernels/Softmax.ts b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
index bd2de9abb08..7209cdb8c1d 100644
--- a/tfjs-core/src/backends/cpu/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
@@ -21,7 +21,7 @@ import {parseAxisParam, sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
-import {div} from './div_impl';
+import {div} from './Div';
 import {exp} from './exp_impl';
 import {max} from './max_impl';
 import {sub} from './sub_impl';
diff --git a/tfjs-core/src/backends/cpu/kernels/div_impl.ts b/tfjs-core/src/backends/cpu/kernels/div_impl.ts
deleted file mode 100644
index dd5e3453054..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/div_impl.ts
+++ /dev/null
@@ -1,20 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {createBinaryOp} from '../utils/kernel_utils';
-
-export const div = createBinaryOp((a: number, b: number) => a / b);
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index e87aa5e8a2a..f06b7717516 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,13 +15,13 @@
  * =============================================================================
  */
 
-// import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
-import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
+// import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
+import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', WEBGL_ENVS, () => {
+describeWithFlags('softmax', CPU_ENVS, () => {
   fit('regular test', async () => {
     console.log('TESTINGGGG');
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));

From 5eb5143c322a83529a2d8913fe94ff84380257a8 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 10:06:07 -0400
Subject: [PATCH 79/84] simplify

---
 tfjs-core/src/backends/cpu/kernels/Exp.ts     | 12 +++++-
 tfjs-core/src/backends/cpu/kernels/Max.ts     | 23 ++++++++++-
 tfjs-core/src/backends/cpu/kernels/Softmax.ts |  8 ++--
 tfjs-core/src/backends/cpu/kernels/Sub.ts     |  3 +-
 tfjs-core/src/backends/cpu/kernels/Sum.ts     | 21 +++++++++-
 .../src/backends/cpu/kernels/exp_impl.ts      | 29 --------------
 .../src/backends/cpu/kernels/max_impl.ts      | 40 -------------------
 .../src/backends/cpu/kernels/sub_impl.ts      | 20 ----------
 .../src/backends/cpu/kernels/sum_impl.ts      | 38 ------------------
 9 files changed, 59 insertions(+), 135 deletions(-)
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/exp_impl.ts
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/max_impl.ts
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/sub_impl.ts
 delete mode 100644 tfjs-core/src/backends/cpu/kernels/sum_impl.ts

diff --git a/tfjs-core/src/backends/cpu/kernels/Exp.ts b/tfjs-core/src/backends/cpu/kernels/Exp.ts
index e607dc21938..a6d40fdd1c4 100644
--- a/tfjs-core/src/backends/cpu/kernels/Exp.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Exp.ts
@@ -17,9 +17,19 @@
 
 import {Exp, ExpInputs} from '../../../kernel_names';
 import {KernelConfig} from '../../../kernel_registry';
+import {TypedArray} from '../../../types';
+import * as util from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 
-import {exp} from './exp_impl';
+export const exp = (x: TypedArray): TypedArray => {
+  const outValues = util.getTypedArrayFromDType('float32', x.length);
+
+  for (let i = 0; i < x.length; ++i) {
+    outValues[i] = Math.exp(x[i]);
+  }
+
+  return outValues;
+};
 
 export const expConfig: KernelConfig = {
   kernelName: Exp,
diff --git a/tfjs-core/src/backends/cpu/kernels/Max.ts b/tfjs-core/src/backends/cpu/kernels/Max.ts
index 1ac73aa4f85..23b018ae9f6 100644
--- a/tfjs-core/src/backends/cpu/kernels/Max.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Max.ts
@@ -18,11 +18,32 @@
 import {Max, MaxAttrs, MaxInputs} from '../../../kernel_names';
 import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
+import {DataType, NumericDataType, TypedArray} from '../../../types';
+import * as util from '../../../util';
 import {sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
-import {max} from './max_impl';
+export const max =
+    (x: TypedArray, reduceSize: number, outShape: number[], dtype: DataType):
+        TypedArray => {
+          const outValues = util.getTypedArrayFromDType(
+              dtype as NumericDataType, util.sizeFromShape(outShape));
+
+          for (let i = 0; i < x.length; ++i) {
+            const offset = i * reduceSize;
+            let max = x[offset];
+            for (let j = 0; j < reduceSize; ++j) {
+              const value = x[offset + j];
+              if (value > max) {
+                max = value;
+              }
+            }
+            outValues[i] = max;
+          }
+
+          return outValues;
+        };
 
 export const maxConfig: KernelConfig = {
   kernelName: Max,
diff --git a/tfjs-core/src/backends/cpu/kernels/Softmax.ts b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
index 7209cdb8c1d..d111b1db746 100644
--- a/tfjs-core/src/backends/cpu/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
@@ -22,10 +22,10 @@ import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
 import {div} from './Div';
-import {exp} from './exp_impl';
-import {max} from './max_impl';
-import {sub} from './sub_impl';
-import {sum} from './sum_impl';
+import {exp} from './Exp';
+import {max} from './Max';
+import {sub} from './Sub';
+import {sum} from './Sum';
 
 interface SoftmaxInputs extends NamedTensorInfoMap {
   x: TensorInfo;
diff --git a/tfjs-core/src/backends/cpu/kernels/Sub.ts b/tfjs-core/src/backends/cpu/kernels/Sub.ts
index 96ed5e074fb..f37ebde371f 100644
--- a/tfjs-core/src/backends/cpu/kernels/Sub.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Sub.ts
@@ -17,7 +17,8 @@
 
 import {Sub} from '../../../kernel_names';
 import {createBinaryKernelConfig} from '../utils/kernel_utils';
+import {createBinaryOp} from '../utils/kernel_utils';
 
-import {sub} from './sub_impl';
+export const sub = createBinaryOp((a: number, b: number) => a - b);
 
 export const subConfig = createBinaryKernelConfig(Sub, sub);
diff --git a/tfjs-core/src/backends/cpu/kernels/Sum.ts b/tfjs-core/src/backends/cpu/kernels/Sum.ts
index 43c4930e8f9..8779e9efb1d 100644
--- a/tfjs-core/src/backends/cpu/kernels/Sum.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Sum.ts
@@ -19,11 +19,30 @@ import {Sum, SumAttrs, SumInputs} from '../../../kernel_names';
 import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
 import {upcastType} from '../../../types';
+import {DataType, NumericDataType, TypedArray} from '../../../types';
+import * as util from '../../../util';
 import {sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
 import {assertNotComplex} from '../cpu_util';
 
-import {sum} from './sum_impl';
+export const sum =
+    (x: TypedArray, reduceSize: number, outShape: number[], dtype: DataType):
+        TypedArray => {
+          const outValues = util.getTypedArrayFromDType(
+              dtype as NumericDataType, util.sizeFromShape(outShape));
+
+          for (let i = 0; i < x.length; ++i) {
+            const offset = i * reduceSize;
+            let sum = 0;
+            for (let j = 0; j < reduceSize; ++j) {
+              const value = x[offset + j];
+              sum += value;
+            }
+            outValues[i] = sum;
+          }
+
+          return outValues;
+        };
 
 export const sumConfig: KernelConfig = {
   kernelName: Sum,
diff --git a/tfjs-core/src/backends/cpu/kernels/exp_impl.ts b/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
deleted file mode 100644
index 111ad85455e..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/exp_impl.ts
+++ /dev/null
@@ -1,29 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {TypedArray} from '../../../types';
-import * as util from '../../../util';
-
-export const exp = (x: TypedArray): TypedArray => {
-  const outValues = util.getTypedArrayFromDType('float32', x.length);
-
-  for (let i = 0; i < x.length; ++i) {
-    outValues[i] = Math.exp(x[i]);
-  }
-
-  return outValues;
-};
diff --git a/tfjs-core/src/backends/cpu/kernels/max_impl.ts b/tfjs-core/src/backends/cpu/kernels/max_impl.ts
deleted file mode 100644
index f03a7a7c8d8..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/max_impl.ts
+++ /dev/null
@@ -1,40 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {DataType, NumericDataType, TypedArray} from '../../../types';
-import * as util from '../../../util';
-
-export const max =
-    (x: TypedArray, reduceSize: number, outShape: number[], dtype: DataType):
-        TypedArray => {
-          const outValues = util.getTypedArrayFromDType(
-              dtype as NumericDataType, util.sizeFromShape(outShape));
-
-          for (let i = 0; i < x.length; ++i) {
-            const offset = i * reduceSize;
-            let max = x[offset];
-            for (let j = 0; j < reduceSize; ++j) {
-              const value = x[offset + j];
-              if (value > max) {
-                max = value;
-              }
-            }
-            outValues[i] = max;
-          }
-
-          return outValues;
-        };
diff --git a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts b/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
deleted file mode 100644
index 78bfa4b40a0..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/sub_impl.ts
+++ /dev/null
@@ -1,20 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {createBinaryOp} from '../utils/kernel_utils';
-
-export const sub = createBinaryOp((a: number, b: number) => a - b);
diff --git a/tfjs-core/src/backends/cpu/kernels/sum_impl.ts b/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
deleted file mode 100644
index 36bcdce8c48..00000000000
--- a/tfjs-core/src/backends/cpu/kernels/sum_impl.ts
+++ /dev/null
@@ -1,38 +0,0 @@
-/**
- * @license
- * Copyright 2020 Google LLC. 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 {DataType, NumericDataType, TypedArray} from '../../../types';
-import * as util from '../../../util';
-
-export const sum =
-    (x: TypedArray, reduceSize: number, outShape: number[], dtype: DataType):
-        TypedArray => {
-          const outValues = util.getTypedArrayFromDType(
-              dtype as NumericDataType, util.sizeFromShape(outShape));
-
-          for (let i = 0; i < x.length; ++i) {
-            const offset = i * reduceSize;
-            let sum = 0;
-            for (let j = 0; j < reduceSize; ++j) {
-              const value = x[offset + j];
-              sum += value;
-            }
-            outValues[i] = sum;
-          }
-
-          return outValues;
-        };

From c594bc72f06a665267fa968c65ef4f8159e8c493 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 12:12:05 -0400
Subject: [PATCH 80/84] div

---
 tfjs-core/src/backends/webgl/backend_webgl.ts | 29 ++++++++++---------
 tfjs-core/src/backends/webgl/kernels/Div.ts   | 14 +++++++--
 .../src/backends/webgl/kernels/Softmax.ts     |  8 ++++-
 tfjs-core/src/ops/softmax_test.ts             |  8 ++---
 4 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 9495a223934..af8cfcb658e 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -1372,17 +1372,17 @@ export class MathBackendWebGL extends KernelBackend {
     return this.reduce(a2D, 'any', a2D.dtype).reshape(outShape);
   }
 
-  realDivide(a: Tensor, b: Tensor): Tensor {
-    const op = binaryop_gpu.DIV;
-    const outputDtype = 'float32';
-    if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
-      const checkOutOfBounds = true;
-      return this.packedBinaryOp(
-          a, b, binaryop_packed_gpu.DIV, outputDtype, checkOutOfBounds);
-    }
-    const program = new BinaryOpProgram(op, a.shape, b.shape);
-    return this.compileAndRun<Tensor>(program, [a, b], outputDtype);
-  }
+  // realDivide(a: Tensor, b: Tensor): Tensor {
+  //   const op = binaryop_gpu.DIV;
+  //   const outputDtype = 'float32';
+  //   if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
+  //     const checkOutOfBounds = true;
+  //     return this.packedBinaryOp(
+  //         a, b, binaryop_packed_gpu.DIV, outputDtype, checkOutOfBounds);
+  //   }
+  //   const program = new BinaryOpProgram(op, a.shape, b.shape);
+  //   return this.compileAndRun<Tensor>(program, [a, b], outputDtype);
+  // }
 
   floorDiv(a: Tensor, b: Tensor): Tensor {
     const op = binaryop_gpu.INT_DIV;
@@ -2487,8 +2487,11 @@ export class MathBackendWebGL extends KernelBackend {
   runWebGLProgram(
       program: GPGPUProgram, inputs: TensorInfo[], outputDtype: DataType,
       customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => void,
-      preventEagerUnpackingOfOutput = false): TensorInfo {
-    const output = this.makeTensorInfo(program.outputShape, outputDtype);
+      preventEagerUnpackingOfOutput = false, output?: TensorInfo): TensorInfo {
+    if (output == null) {
+      output = this.makeTensorInfo(program.outputShape, outputDtype);
+    }
+
     const outData = this.texData.get(output.dataId);
     if (program.packedOutput) {
       outData.isPacked = true;
diff --git a/tfjs-core/src/backends/webgl/kernels/Div.ts b/tfjs-core/src/backends/webgl/kernels/Div.ts
index 2be55ee8966..e7aec3784ae 100644
--- a/tfjs-core/src/backends/webgl/kernels/Div.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Div.ts
@@ -15,6 +15,7 @@
  * =============================================================================
  */
 
+import {backend_util} from '../../..';
 import {env} from '../../../environment';
 import {NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 import {MathBackendWebGL} from '../backend_webgl';
@@ -29,13 +30,16 @@ interface DivInputs extends NamedTensorInfoMap {
 }
 
 export const divImpl =
-    (a: TensorInfo, b: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
+    (a: TensorInfo, b: TensorInfo, outTensorInfo: TensorInfo,
+     backend: MathBackendWebGL): TensorInfo => {
       let program = new BinaryOpProgram(binaryop_gpu.DIV, a.shape, b.shape);
       if (env().getBool('WEBGL_PACK_BINARY_OPERATIONS')) {
         program = new BinaryOpPackedProgram(
             binaryop_packed_gpu.DIV, a.shape, b.shape, true);
       }
-      const output = backend.runWebGLProgram(program, [a, b], 'float32');
+      const output = backend.runWebGLProgram(
+          program, [a, b], 'float32', null /* custom setup */,
+          false /* prevent eager unpacking of output */, outTensorInfo);
       return output;
     };
 
@@ -45,7 +49,11 @@ registerKernel({
   kernelFunc: ({inputs, backend}) => {
     const {a, b} = inputs as DivInputs;
     const webglBackend = backend as MathBackendWebGL;
-    const out = divImpl(a, b, webglBackend);
+
+    const outShape = backend_util.assertAndGetBroadcastShape(a.shape, b.shape);
+    const outTensorInfo = webglBackend.makeTensorInfo(outShape, a.dtype);
+
+    const out = divImpl(a, b, outTensorInfo, webglBackend);
 
     return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
   }
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index e286a623de0..e4c4e43294d 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -15,6 +15,7 @@
  * =============================================================================
  */
 
+import {backend_util} from '../../..';
 import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
 // import {parseAxisParam, sizeFromShape} from '../../../util';
@@ -55,7 +56,12 @@ registerKernel({
 
     const summed = sumImpl(exponentiated, reduceShape, outShape, webglBackend);
 
-    const out = divImpl(exponentiated, summed, webglBackend);
+    const divOutTensorInfo = webglBackend.makeTensorInfo(
+        backend_util.assertAndGetBroadcastShape(
+            exponentiated.shape, summed.shape),
+        exponentiated.dtype);
+
+    const out = divImpl(exponentiated, summed, divOutTensorInfo, webglBackend);
 
     webglBackend.disposeData(max.dataId);
     webglBackend.disposeData(subtracted.dataId);
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index f06b7717516..652b54fbb4a 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -15,14 +15,14 @@
  * =============================================================================
  */
 
-// import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
-import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
+import {WEBGL_ENVS} from '../backends/webgl/backend_webgl_test_registry';
+// import {CPU_ENVS} from '../backends/cpu/backend_cpu_test_registry';
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('softmax', CPU_ENVS, () => {
-  fit('regular test', async () => {
+describeWithFlags('softmax', WEBGL_ENVS, () => {
+  it('regular test', async () => {
     console.log('TESTINGGGG');
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 

From 8b4ca3d8f61edeac35f025a64fb35cdefcf26fbc Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 12:26:33 -0400
Subject: [PATCH 81/84] add arg

---
 tfjs-core/src/backends/webgl/backend_webgl.ts | 19 +++++++++++--------
 tfjs-core/src/backends/webgl/kernels/Div.ts   |  5 ++---
 tfjs-core/src/backends/webgl/reshape.ts       |  4 ++--
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index af8cfcb658e..6048657d896 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -2440,8 +2440,8 @@ export class MathBackendWebGL extends KernelBackend {
     const program = new PackProgram(input.shape);
     const preventEagerUnpackingOutput = true;
     return this.runWebGLProgram(
-        program, [input], input.dtype, null /* customSetup */,
-        preventEagerUnpackingOutput);
+        program, [input], input.dtype, null /* out info */,
+        null /* customSetup */, preventEagerUnpackingOutput);
   }
 
   private packedReshape(input: TensorInfo, afterShape: number[]): TensorInfo {
@@ -2461,8 +2461,8 @@ export class MathBackendWebGL extends KernelBackend {
     const program = new ReshapePackedProgram(afterShapeAs3D, input3DShape);
     const preventEagerUnpackingOfOutput = true;
     const output = this.runWebGLProgram(
-        program, [input3D], input.dtype, null /* customSetup */,
-        preventEagerUnpackingOfOutput);
+        program, [input3D], input.dtype, null /* out info */,
+        null /* customSetup */, preventEagerUnpackingOfOutput);
     return {dataId: output.dataId, shape: afterShape, dtype: output.dtype};
   }
 
@@ -2480,14 +2480,16 @@ export class MathBackendWebGL extends KernelBackend {
     const preventEagerUnpackingOfOutput = true;
     const out = this.runWebGLProgram(
         program, [{shape: shapeAs3D, dtype, dataId}], dtype,
-        null /* customSetup */, preventEagerUnpackingOfOutput);
+        null /* out info */, null /* customSetup */,
+        preventEagerUnpackingOfOutput);
     return {dtype, shape, dataId: out.dataId};
   }
 
   runWebGLProgram(
       program: GPGPUProgram, inputs: TensorInfo[], outputDtype: DataType,
+      output?: TensorInfo,
       customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => void,
-      preventEagerUnpackingOfOutput = false, output?: TensorInfo): TensorInfo {
+      preventEagerUnpackingOfOutput = false): TensorInfo {
     if (output == null) {
       output = this.makeTensorInfo(program.outputShape, outputDtype);
     }
@@ -2618,7 +2620,7 @@ export class MathBackendWebGL extends KernelBackend {
       preventEagerUnpackingOfOutput = false): K {
     outputDtype = outputDtype || inputs[0].dtype;
     const outInfo = this.runWebGLProgram(
-        program, inputs, outputDtype, customSetup,
+        program, inputs, outputDtype, null /* out info */, customSetup,
         preventEagerUnpackingOfOutput);
     return ENGINE.makeTensorFromDataId(
                outInfo.dataId, outInfo.shape, outInfo.dtype) as {} as K;
@@ -2745,7 +2747,8 @@ export class MathBackendWebGL extends KernelBackend {
       // WEBGL_PACK.
       const preventEagerUnpacking = true;
       const encodedOutputTarget = this.runWebGLProgram(
-          program, [tempDenseInputHandle], dtype, null, preventEagerUnpacking);
+          program, [tempDenseInputHandle], dtype, null /* out info */, null,
+          preventEagerUnpacking);
 
       // Have the original texture assume the identity of the encoded output.
       const outputTexData = this.texData.get(encodedOutputTarget.dataId);
diff --git a/tfjs-core/src/backends/webgl/kernels/Div.ts b/tfjs-core/src/backends/webgl/kernels/Div.ts
index e7aec3784ae..ea6eb839735 100644
--- a/tfjs-core/src/backends/webgl/kernels/Div.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Div.ts
@@ -37,9 +37,8 @@ export const divImpl =
         program = new BinaryOpPackedProgram(
             binaryop_packed_gpu.DIV, a.shape, b.shape, true);
       }
-      const output = backend.runWebGLProgram(
-          program, [a, b], 'float32', null /* custom setup */,
-          false /* prevent eager unpacking of output */, outTensorInfo);
+      const output =
+          backend.runWebGLProgram(program, [a, b], 'float32', outTensorInfo);
       return output;
     };
 
diff --git a/tfjs-core/src/backends/webgl/reshape.ts b/tfjs-core/src/backends/webgl/reshape.ts
index 2b3e47b3ef0..cd69d7e8427 100644
--- a/tfjs-core/src/backends/webgl/reshape.ts
+++ b/tfjs-core/src/backends/webgl/reshape.ts
@@ -42,8 +42,8 @@ const packedReshape =
       const program = new ReshapePackedProgram(afterShapeAs3D, input3DShape);
       const preventEagerUnpackingOfOutput = true;
       const output = backend.runWebGLProgram(
-          program, [input3D], input.dtype, null /* customSetup */,
-          preventEagerUnpackingOfOutput);
+          program, [input3D], input.dtype, null /* out info */,
+          null /* customSetup */, preventEagerUnpackingOfOutput);
       return {dataId: output.dataId, shape: afterShape, dtype: output.dtype};
     };
 

From a332b15be824e15e4d03f07f9fec38fabca1cf5f Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 12:50:45 -0400
Subject: [PATCH 82/84] reduce

---
 tfjs-core/src/backends/webgl/kernels/Max.ts     |  8 +++++---
 tfjs-core/src/backends/webgl/kernels/Softmax.ts | 10 ++++++++--
 tfjs-core/src/backends/webgl/kernels/Sum.ts     |  9 +++++----
 tfjs-core/src/backends/webgl/reduce.ts          |  6 +++---
 4 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Max.ts b/tfjs-core/src/backends/webgl/kernels/Max.ts
index 4f0e546b544..5fa8d5b3fd4 100644
--- a/tfjs-core/src/backends/webgl/kernels/Max.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Max.ts
@@ -32,7 +32,7 @@ interface MaxAttrs extends NamedAttrMap {
 
 export const maxImpl =
     (x: TensorInfo, reduceShape: number[], outShape: number[],
-     backend: MathBackendWebGL): TensorInfo => {
+     outInfo: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
       const inSize = sizeFromShape(reduceShape);
       const xSize = sizeFromShape(x.shape);
       const batchSize = xSize / inSize;
@@ -40,7 +40,7 @@ export const maxImpl =
       x = reshape(x, [batchSize, inSize], backend);
 
       return reshape(
-          reduce(x, reduceShape, x.dtype, backend), outShape, backend);
+          reduce(x, reduceShape, x.dtype, outInfo, backend), outShape, backend);
     };
 
 registerKernel({
@@ -56,7 +56,9 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(x.shape, axes);
 
-    const out = maxImpl(x, reduceShape, outShape, webglBackend);
+    const outTensorInfo = webglBackend.makeTensorInfo(outShape, x.dtype);
+
+    const out = maxImpl(x, reduceShape, outShape, outTensorInfo, webglBackend);
 
     return {dataId: out.dataId, shape: outShape, dtype: x.dtype};
   }
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index e4c4e43294d..e181022d072 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -49,12 +49,18 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
 
-    const max = maxImpl(logits, reduceShape, outShape, webglBackend);
+    const maxOutTensorInfo =
+        webglBackend.makeTensorInfo(outShape, logits.dtype);
+    const max =
+        maxImpl(logits, reduceShape, outShape, maxOutTensorInfo, webglBackend);
 
     const subtracted = subImpl(logits, max, webglBackend);
     const exponentiated = expImpl(subtracted, webglBackend);
 
-    const summed = sumImpl(exponentiated, reduceShape, outShape, webglBackend);
+    const sumOutTensorInfo =
+        webglBackend.makeTensorInfo(outShape, logits.dtype);
+    const summed = sumImpl(
+        exponentiated, reduceShape, outShape, sumOutTensorInfo, webglBackend);
 
     const divOutTensorInfo = webglBackend.makeTensorInfo(
         backend_util.assertAndGetBroadcastShape(
diff --git a/tfjs-core/src/backends/webgl/kernels/Sum.ts b/tfjs-core/src/backends/webgl/kernels/Sum.ts
index 7f51ec5f1a6..998d63fc521 100644
--- a/tfjs-core/src/backends/webgl/kernels/Sum.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Sum.ts
@@ -34,7 +34,7 @@ interface SumAttrs extends NamedAttrMap {
 
 export const sumImpl =
     (x: TensorInfo, reduceShape: number[], outShape: number[],
-     backend: MathBackendWebGL): TensorInfo => {
+     outInfo: TensorInfo, backend: MathBackendWebGL): TensorInfo => {
       const inSize = sizeFromShape(reduceShape);
       const xSize = sizeFromShape(x.shape);
       const batchSize = xSize / inSize;
@@ -42,8 +42,8 @@ export const sumImpl =
       x = reshape(x, [batchSize, inSize], backend);
 
       return reshape(
-          reduce(x, reduceShape, sumOutType(x.dtype), backend), outShape,
-          backend);
+          reduce(x, reduceShape, sumOutType(x.dtype), outInfo, backend),
+          outShape, backend);
     };
 
 registerKernel({
@@ -59,7 +59,8 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(x.shape, axes);
 
-    const out = sumImpl(x, reduceShape, outShape, webglBackend);
+    const outTensorInfo = webglBackend.makeTensorInfo(outShape, x.dtype);
+    const out = sumImpl(x, reduceShape, outShape, outTensorInfo, webglBackend);
 
     webglBackend.disposeData(x);
 
diff --git a/tfjs-core/src/backends/webgl/reduce.ts b/tfjs-core/src/backends/webgl/reduce.ts
index 4cee37841e2..c63587e0af6 100644
--- a/tfjs-core/src/backends/webgl/reduce.ts
+++ b/tfjs-core/src/backends/webgl/reduce.ts
@@ -23,13 +23,13 @@ import {MathBackendWebGL} from './backend_webgl';
 import {ReduceProgram} from './reduce_gpu';
 
 export const reduce =
-    (x: TensorInfo, reduceShape: number[], dtype: DataType,
+    (x: TensorInfo, reduceShape: number[], dtype: DataType, outInfo: TensorInfo,
      backend: MathBackendWebGL): TensorInfo => {
       const [batchSize, inSize] = x.shape;
       const windowSize = computeOptimalWindowSize(inSize);
       const reduceInfo = {windowSize, inSize, batchSize};
       const program = new ReduceProgram(reduceInfo, 'sum');
-      const output = backend.runWebGLProgram(program, [x], dtype);
+      const output = backend.runWebGLProgram(program, [x], dtype, outInfo);
 
       backend.disposeData(x.dataId);
 
@@ -37,5 +37,5 @@ export const reduce =
         return output;
       }
 
-      return reduce(output, reduceShape, dtype, backend);
+      return reduce(output, reduceShape, dtype, outInfo, backend);
     };

From ea195f616600e55055bb460d38ee07d9e537aec5 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 13:53:08 -0400
Subject: [PATCH 83/84] updates

---
 tfjs-core/src/backends/cpu/kernels/Softmax.ts | 19 ++----
 .../src/backends/cpu/register_all_kernels.ts  |  4 +-
 tfjs-core/src/backends/webgl/backend_webgl.ts |  1 -
 .../src/backends/webgl/kernels/Softmax.ts     | 64 +++++++++----------
 tfjs-core/src/backends/webgl/reduce.ts        |  8 ++-
 .../backends/webgl/register_all_kernels.ts    |  2 +
 .../src/backends/webgl/webgl_ops_test.ts      |  2 +-
 tfjs-core/src/engine.ts                       |  4 --
 tfjs-core/src/kernel_names.ts                 |  6 ++
 tfjs-core/src/ops/softmax_test.ts             |  7 +-
 10 files changed, 60 insertions(+), 57 deletions(-)

diff --git a/tfjs-core/src/backends/cpu/kernels/Softmax.ts b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
index d111b1db746..06731f62f2a 100644
--- a/tfjs-core/src/backends/cpu/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/cpu/kernels/Softmax.ts
@@ -15,7 +15,8 @@
  * =============================================================================
  */
 
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {Softmax, SoftmaxAttrs, SoftmaxInputs} from '../../../kernel_names';
+import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
 import {parseAxisParam, sizeFromShape} from '../../../util';
 import {MathBackendCPU} from '../backend_cpu';
@@ -27,20 +28,12 @@ import {max} from './Max';
 import {sub} from './Sub';
 import {sum} from './Sum';
 
-interface SoftmaxInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface SoftmaxAttrs extends NamedAttrMap {
-  dim: number;
-}
-
-registerKernel({
-  kernelName: 'Softmax',
+export const softmaxConfig: KernelConfig = {
+  kernelName: Softmax,
   backendName: 'cpu',
   kernelFunc: ({inputs, attrs, backend}) => {
     const {logits} = inputs as SoftmaxInputs;
-    const {dim} = attrs as SoftmaxAttrs;
+    const {dim} = attrs as {} as SoftmaxAttrs;
     const cpuBackend = backend as MathBackendCPU;
     assertNotComplex(logits, 'softmax');
 
@@ -68,4 +61,4 @@ registerKernel({
     const dataId = cpuBackend.write(resultData, resultShape, logits.dtype);
     return {dataId, shape: resultShape, dtype: logits.dtype};
   }
-});
+};
diff --git a/tfjs-core/src/backends/cpu/register_all_kernels.ts b/tfjs-core/src/backends/cpu/register_all_kernels.ts
index 4d7f2e983cf..ef9282a38a2 100644
--- a/tfjs-core/src/backends/cpu/register_all_kernels.ts
+++ b/tfjs-core/src/backends/cpu/register_all_kernels.ts
@@ -21,12 +21,14 @@ import {KernelConfig, registerKernel} from '../../kernel_registry';
 
 import {divConfig} from './kernels/Div';
 import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5';
+import {softmaxConfig} from './kernels/Softmax';
 import {squareConfig} from './kernels/Square';
 import {squaredDifferenceConfig} from './kernels/SquaredDifference';
 
 // List all kernel configs here
 const kernelConfigs: KernelConfig[] = [
-  nonMaxSuppressionV5Config, squareConfig, squaredDifferenceConfig, divConfig
+  nonMaxSuppressionV5Config, squareConfig, squaredDifferenceConfig, divConfig,
+  softmaxConfig
 ];
 
 for (const kernelConfig of kernelConfigs) {
diff --git a/tfjs-core/src/backends/webgl/backend_webgl.ts b/tfjs-core/src/backends/webgl/backend_webgl.ts
index 6048657d896..321e693a55a 100644
--- a/tfjs-core/src/backends/webgl/backend_webgl.ts
+++ b/tfjs-core/src/backends/webgl/backend_webgl.ts
@@ -2701,7 +2701,6 @@ export class MathBackendWebGL extends KernelBackend {
       // Array is already on GPU. No-op.
       return;
     }
-    console.log('UPLOAD TO GPU');
     const shouldTimeProgram = this.activeTimers != null;
     let start: number;
     if (shouldTimeProgram) {
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index e181022d072..3dc504888bd 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -15,33 +15,26 @@
  * =============================================================================
  */
 
-import {backend_util} from '../../..';
-import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
+import {Softmax, SoftmaxAttrs, SoftmaxInputs} from '../../../kernel_names';
+// import {backend_util} from '../../..';
+import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
 // import {parseAxisParam, sizeFromShape} from '../../../util';
 import {parseAxisParam} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
 
-import {divImpl} from './Div';
-import {expImpl} from './Exp';
+// import {divImpl} from './Div';
+// import {expImpl} from './Exp';
 import {maxImpl} from './Max';
-import {subImpl} from './Sub';
-import {sumImpl} from './Sum';
+// import {subImpl} from './Sub';
+// import {sumImpl} from './Sum';
 
-interface SoftmaxInputs extends NamedTensorInfoMap {
-  x: TensorInfo;
-}
-
-interface SoftmaxAttrs extends NamedAttrMap {
-  dim: number;
-}
-
-registerKernel({
-  kernelName: 'Softmax',
+export const softmaxConfig: KernelConfig = {
+  kernelName: Softmax,
   backendName: 'webgl',
   kernelFunc: ({inputs, attrs, backend}) => {
     const {logits} = inputs as SoftmaxInputs;
-    const {dim} = attrs as SoftmaxAttrs;
+    const {dim} = attrs as {} as SoftmaxAttrs;
     const webglBackend = backend as MathBackendWebGL;
 
     const axes = parseAxisParam([dim], logits.shape);
@@ -54,26 +47,29 @@ registerKernel({
     const max =
         maxImpl(logits, reduceShape, outShape, maxOutTensorInfo, webglBackend);
 
-    const subtracted = subImpl(logits, max, webglBackend);
-    const exponentiated = expImpl(subtracted, webglBackend);
+    // const subtracted = subImpl(logits, max, webglBackend);
+    // const exponentiated = expImpl(subtracted, webglBackend);
 
-    const sumOutTensorInfo =
-        webglBackend.makeTensorInfo(outShape, logits.dtype);
-    const summed = sumImpl(
-        exponentiated, reduceShape, outShape, sumOutTensorInfo, webglBackend);
+    // const sumOutTensorInfo =
+    //     webglBackend.makeTensorInfo(outShape, logits.dtype);
+    // const summed = sumImpl(
+    //     exponentiated, reduceShape, outShape, sumOutTensorInfo,
+    //     webglBackend);
 
-    const divOutTensorInfo = webglBackend.makeTensorInfo(
-        backend_util.assertAndGetBroadcastShape(
-            exponentiated.shape, summed.shape),
-        exponentiated.dtype);
+    // const divOutTensorInfo = webglBackend.makeTensorInfo(
+    //     backend_util.assertAndGetBroadcastShape(
+    //         exponentiated.shape, summed.shape),
+    //     exponentiated.dtype);
 
-    const out = divImpl(exponentiated, summed, divOutTensorInfo, webglBackend);
+    // const out = divImpl(exponentiated, summed, divOutTensorInfo,
+    // webglBackend);
 
-    webglBackend.disposeData(max.dataId);
-    webglBackend.disposeData(subtracted.dataId);
-    webglBackend.disposeData(exponentiated.dataId);
-    webglBackend.disposeData(summed.dataId);
+    // webglBackend.disposeData(max.dataId);
+    // webglBackend.disposeData(subtracted.dataId);
+    // webglBackend.disposeData(exponentiated.dataId);
+    // webglBackend.disposeData(summed.dataId);
 
-    return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
+    // return {dataId: out.dataId, shape: out.shape, dtype: out.dtype};
+    return {dataId: max.dataId, shape: max.shape, dtype: max.dtype};
   }
-});
+};
diff --git a/tfjs-core/src/backends/webgl/reduce.ts b/tfjs-core/src/backends/webgl/reduce.ts
index c63587e0af6..97fa91a3ee5 100644
--- a/tfjs-core/src/backends/webgl/reduce.ts
+++ b/tfjs-core/src/backends/webgl/reduce.ts
@@ -37,5 +37,11 @@ export const reduce =
         return output;
       }
 
-      return reduce(output, reduceShape, dtype, outInfo, backend);
+      const [newBatchSize, newInSize] = output.shape;
+      const newWindowSize = computeOptimalWindowSize(newInSize);
+      const newOutShape = [newBatchSize, Math.ceil(newInSize / newWindowSize)];
+
+      const newOutInfo = backend.makeTensorInfo(newOutShape, output.dtype);
+
+      return reduce(output, reduceShape, dtype, newOutInfo, backend);
     };
diff --git a/tfjs-core/src/backends/webgl/register_all_kernels.ts b/tfjs-core/src/backends/webgl/register_all_kernels.ts
index f7913ac6b21..8521e55d446 100644
--- a/tfjs-core/src/backends/webgl/register_all_kernels.ts
+++ b/tfjs-core/src/backends/webgl/register_all_kernels.ts
@@ -18,6 +18,7 @@ import {KernelConfig, registerKernel} from '../../kernel_registry';
 
 import {fromPixelsConfig} from './kernels/FromPixels';
 import {nonMaxSuppressionV5Config} from './kernels/NonMaxSuppressionV5';
+import {softmaxConfig} from './kernels/Softmax';
 import {squareConfig} from './kernels/Square';
 import {squaredDifferenceConfig} from './kernels/SquaredDifference';
 
@@ -27,6 +28,7 @@ const kernelConfigs: KernelConfig[] = [
   nonMaxSuppressionV5Config,
   squareConfig,
   squaredDifferenceConfig,
+  softmaxConfig,
 ];
 
 for (const kernelConfig of kernelConfigs) {
diff --git a/tfjs-core/src/backends/webgl/webgl_ops_test.ts b/tfjs-core/src/backends/webgl/webgl_ops_test.ts
index e5db4bb177c..3dc5c1ebab1 100644
--- a/tfjs-core/src/backends/webgl/webgl_ops_test.ts
+++ b/tfjs-core/src/backends/webgl/webgl_ops_test.ts
@@ -444,7 +444,7 @@ describeWithFlags('gramSchmidt-non-tiny', WEBGL_ENVS, () => {
 });
 
 describeWithFlags('matmul webgl-only', WEBGL_ENVS, () => {
-  fit('Matrix times vector, large matrix', async () => {
+  it('Matrix times vector, large matrix', async () => {
     const maxTexSize = 16000;
     const sharedDim = maxTexSize + 4;
     const matrix = tf.buffer<Rank.R2>([2, sharedDim], 'float32');
diff --git a/tfjs-core/src/engine.ts b/tfjs-core/src/engine.ts
index ac2b6c40473..2f7c17c4e18 100644
--- a/tfjs-core/src/engine.ts
+++ b/tfjs-core/src/engine.ts
@@ -522,7 +522,6 @@ export class Engine implements TensorTracker, DataMover {
     const dataIdsLeaked =
         numDataIdsAfter - numDataIdsBefore - numOutputDataIds - numMoves;
 
-    console.log(numDataIdsAfter, numDataIdsBefore, numOutputDataIds, numMoves);
     if (dataIdsLeaked > 0) {
       throw new Error(
           `Backend '${this.backendName}' has an internal memory leak ` +
@@ -568,13 +567,10 @@ export class Engine implements TensorTracker, DataMover {
     let out: TensorInfo|TensorInfo[];
     if (kernel != null) {
       kernelFunc = () => {
-        console.log('STARTING TO TRACK');
         const numDataIdsBefore = this.backend.numDataIds();
         out = kernel.kernelFunc({inputs, attrs, backend: this.backend});
         const outInfos = Array.isArray(out) ? out : [out];
         if (this.shouldCheckForMemLeaks()) {
-          console.log('CHECK KERNEL', numDataIdsBefore);
-          console.log(outInfos);
           this.checkKernelForMemLeak(kernelName, numDataIdsBefore, outInfos);
         }
         const outTensors = outInfos.map(
diff --git a/tfjs-core/src/kernel_names.ts b/tfjs-core/src/kernel_names.ts
index eedcc24980e..cd764cd0508 100644
--- a/tfjs-core/src/kernel_names.ts
+++ b/tfjs-core/src/kernel_names.ts
@@ -34,6 +34,12 @@ export interface MaxAttrs {
   axes: number[];
 }
 
+export const Softmax = 'Softmax';
+export type SoftmaxInputs = Pick<NamedTensorInfoMap, 'logits'>;
+export interface SoftmaxAttrs {
+  dim: number;
+}
+
 export const SquaredDifference = 'SquaredDifference';
 
 export const Square = 'Square';
diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
index 652b54fbb4a..6cb529425d8 100644
--- a/tfjs-core/src/ops/softmax_test.ts
+++ b/tfjs-core/src/ops/softmax_test.ts
@@ -22,11 +22,14 @@ import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
 describeWithFlags('softmax', WEBGL_ENVS, () => {
-  it('regular test', async () => {
+  fit('regular test', async () => {
     console.log('TESTINGGGG');
     const y = tf.softmax(tf.tensor1d([2, 1, 3]));
 
-    expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
+    const data = await y.data();
+    console.log(Array.from(data));
+
+    // expectArraysClose(await y.data(), [0.24472847, 0.09003057, 0.66524095]);
     // expectArraysClose(await y.sum().data(), 1);
     // const x = tf.tensor1d([1, 2, 3]);
     // x.softmax().print();

From 02f1201b0303ddcc54c7a350b1b592fa60ac1357 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Fri, 20 Mar 2020 16:07:11 -0400
Subject: [PATCH 84/84] reduce shape helper

---
 tfjs-core/src/backends/webgl/kernels/Max.ts     |  4 +++-
 tfjs-core/src/backends/webgl/kernels/Softmax.ts |  6 +++++-
 tfjs-core/src/ops/reduce_util.ts                | 12 +++++++++++-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/tfjs-core/src/backends/webgl/kernels/Max.ts b/tfjs-core/src/backends/webgl/kernels/Max.ts
index 5fa8d5b3fd4..51afc8bc633 100644
--- a/tfjs-core/src/backends/webgl/kernels/Max.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Max.ts
@@ -17,6 +17,7 @@
 
 import {NamedAttrMap, NamedTensorInfoMap, registerKernel, TensorInfo} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
+import {reduceOutShapeFromInShape} from '../../../ops/reduce_util';
 import {sizeFromShape} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
 import {reduce} from '../reduce';
@@ -56,7 +57,8 @@ registerKernel({
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(x.shape, axes);
 
-    const outTensorInfo = webglBackend.makeTensorInfo(outShape, x.dtype);
+    const outTensorInfo = webglBackend.makeTensorInfo(
+        reduceOutShapeFromInShape(x.shape, reduceShape), x.dtype);
 
     const out = maxImpl(x, reduceShape, outShape, outTensorInfo, webglBackend);
 
diff --git a/tfjs-core/src/backends/webgl/kernels/Softmax.ts b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
index 3dc504888bd..abe1fbe697f 100644
--- a/tfjs-core/src/backends/webgl/kernels/Softmax.ts
+++ b/tfjs-core/src/backends/webgl/kernels/Softmax.ts
@@ -19,6 +19,7 @@ import {Softmax, SoftmaxAttrs, SoftmaxInputs} from '../../../kernel_names';
 // import {backend_util} from '../../..';
 import {KernelConfig} from '../../../kernel_registry';
 import * as axis_util from '../../../ops/axis_util';
+import {reduceOutShapeFromInShape} from '../../../ops/reduce_util';
 // import {parseAxisParam, sizeFromShape} from '../../../util';
 import {parseAxisParam} from '../../../util';
 import {MathBackendWebGL} from '../backend_webgl';
@@ -26,6 +27,7 @@ import {MathBackendWebGL} from '../backend_webgl';
 // import {divImpl} from './Div';
 // import {expImpl} from './Exp';
 import {maxImpl} from './Max';
+
 // import {subImpl} from './Sub';
 // import {sumImpl} from './Sum';
 
@@ -42,8 +44,10 @@ export const softmaxConfig: KernelConfig = {
     const [outShape, reduceShape] =
         axis_util.computeOutAndReduceShapes(logits.shape, axes);
 
+    const maxOutShape = reduceOutShapeFromInShape(logits.shape, reduceShape);
+
     const maxOutTensorInfo =
-        webglBackend.makeTensorInfo(outShape, logits.dtype);
+        webglBackend.makeTensorInfo(maxOutShape, logits.dtype);
     const max =
         maxImpl(logits, reduceShape, outShape, maxOutTensorInfo, webglBackend);
 
diff --git a/tfjs-core/src/ops/reduce_util.ts b/tfjs-core/src/ops/reduce_util.ts
index cb9f27d158a..a772aa03372 100644
--- a/tfjs-core/src/ops/reduce_util.ts
+++ b/tfjs-core/src/ops/reduce_util.ts
@@ -19,7 +19,7 @@
  * Inputs of size above this threshold will be parallelized by calling multiple
  * shader programs.
  */
-import {nearestDivisor} from '../util';
+import {nearestDivisor, sizeFromShape} from '../util';
 
 export const PARALLELIZE_THRESHOLD = 30;
 
@@ -35,3 +35,13 @@ export function computeOptimalWindowSize(inSize: number): number {
   }
   return nearestDivisor(inSize, Math.floor(Math.sqrt(inSize)));
 }
+
+export function reduceOutShapeFromInShape(
+    xShape: number[], reduceShape: number[]): number[] {
+  const xSize = sizeFromShape(xShape);
+  const inSize = sizeFromShape(reduceShape);
+  const batchSize = xSize / inSize;
+  const windowSize = computeOptimalWindowSize(inSize);
+  const outSize = Math.ceil(inSize / windowSize);
+  return [batchSize, outSize];
+}