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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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 0024ad6edc9a74532593337be7aa6d808c3a77b5 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 11:47:49 -0500
Subject: [PATCH 40/42] add h

---
 tfjs-backend-wasm/src/cc/BUILD                |  1 +
 tfjs-backend-wasm/src/cc/kernels/Softmax.h    | 29 +++++++++++++++++++
 .../src/cc/kernels/Softmax_test.cc            |  2 +-
 3 files changed, 31 insertions(+), 1 deletion(-)
 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 21117098ccd..0126adfa16e 100644
--- a/tfjs-backend-wasm/src/cc/BUILD
+++ b/tfjs-backend-wasm/src/cc/BUILD
@@ -672,6 +672,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..ee705e4a1ab
--- /dev/null
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.h
@@ -0,0 +1,29 @@
+/* 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.
+ * ===========================================================================*/
+
+#ifndef KERNELS_SOFTMAX_H_
+#define KERNELS_SOFTMAX_H_
+
+#include <cstddef>
+
+namespace tfjs {
+namespace wasm {
+
+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 e010471bdf2..bb8420c6672 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -18,7 +18,7 @@
 #include <vector>
 
 #include "src/cc/backend.h"
-#include "src/cc/util.h"
+#include "src/cc/kernels/Softmax.h"
 
 TEST(SOFTMAX, xnn_operator_lifetime) {
   tfjs::wasm::init();

From 9de07dac794df3e75e6b9f904ad578c00d113344 Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 12:51:12 -0500
Subject: [PATCH 41/42] fix test

---
 tfjs-backend-wasm/src/cc/kernels/Sigmoid.h       | 1 -
 tfjs-backend-wasm/src/cc/kernels/Softmax.cc      | 2 +-
 tfjs-backend-wasm/src/cc/kernels/Softmax.h       | 3 ++-
 tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc | 3 ---
 4 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/tfjs-backend-wasm/src/cc/kernels/Sigmoid.h b/tfjs-backend-wasm/src/cc/kernels/Sigmoid.h
index 4ed39787b34..f24a6415b03 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Sigmoid.h
+++ b/tfjs-backend-wasm/src/cc/kernels/Sigmoid.h
@@ -22,7 +22,6 @@ namespace wasm {
 extern "C" {
 
 void Sigmoid(const size_t x_id, const size_t out_id);
-
 }
 
 }  // namespace wasm
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
index c1cbe6ad3b6..84dc027cb52 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.cc
@@ -21,9 +21,9 @@
 #include <cstddef>
 #include <map>
 #include <tuple>
-#include <vector>
 
 #include "src/cc/backend.h"
+#include "src/cc/kernels/Softmax.h"
 #include "src/cc/util.h"
 
 namespace {
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax.h b/tfjs-backend-wasm/src/cc/kernels/Softmax.h
index ee705e4a1ab..f2cd8ebc033 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax.h
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax.h
@@ -19,9 +19,10 @@
 
 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
diff --git a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
index bb8420c6672..eeec4413617 100644
--- a/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
+++ b/tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc
@@ -45,9 +45,6 @@ TEST(SOFTMAX, xnn_operator_lifetime) {
   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();
-
   // 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 4827669efdc1bdfe30622d888ca61a34d5f5b52b Mon Sep 17 00:00:00 2001
From: Ann Yuan <annyuan@gmail.com>
Date: Thu, 13 Feb 2020 13:02:34 -0500
Subject: [PATCH 42/42] save