[wasm] Add Softmax, Neg, NotEqual.

[annxingyuan]

Changes

• Add Softmax to KernelBackend interface, change softmax op to call kernel directly.
• Added Softmax kernel to WASM backend. Also added neg and notEqual, which are used in certain softmax tests.
• Remove ffast-math from wasm build configuration because it causes XNNPACK to produce incorrect outputs.

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

---

This change is 

[nsthorat]

Run yarn bazel:format to get the CI to pass

Reviewed 3 of 14 files at r1, 2 of 10 files at r2, 15 of 17 files at r3.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, and @nsthorat)

---

tfjs-backend-wasm/.bazelrc, line 16 at r3 (raw file):

build:wasm --cxxopt="-fno-exceptions"
build:wasm --cxxopt="-fomit-frame-pointer"
build:wasm --cxxopt="-ffast-math"

can you add this to the PR description?

---

tfjs-backend-wasm/src/setup_test.ts, line 40 at r3 (raw file):

    include: 'softmax',
    excludes: [
      'gradient'  // Gradient not yet implemented.

do you have to do this? i think all the kernels needed for the gradient are defined in wasm

---

tfjs-core/src/backends/backend.ts, line 308 at r3 (raw file):

    return notYetImplemented('expm1');
  }
  softmax<T extends Tensor>(x: T, dim: number): T {

let's not do this on the backend interface at all, rather we should do this using registerKernel with the string 'Softmax'. Let backend_cpu and webgl have implementations that are modular too.

---

tfjs-core/src/backends/cpu/backend_cpu.ts, line 388 at r3 (raw file):

    const sumExp = this.sum(b, axes).reshape(expandedShape);

    return b.div(sumExp);

no chaining API here and above, use methods directly, e.g. https://github.com/tensorflow/tfjs/blob/master/tfjs-core/src/backends/cpu/backend_cpu.ts#L348

This actually should cause a test failure since we have something that makes sure you don't leave the backend

---

tfjs-core/src/backends/webgl/backend_webgl.ts, line 1600 at r3 (raw file):

  }

  softmax<T extends Tensor>(logits: T, dim: number): T {

here too make this modular

---

tfjs-core/src/ops/softmax.ts, line 61 at r3 (raw file):

  return ENGINE.runKernelFunc(
      (backend, save) => {
        const y = backend.softmax($logits, dim);

you shouldn't have to pass a function here once everyone is in modular land :)

---

tfjs-core/src/ops/softmax.ts, line 75 at r3 (raw file):

        };
      },
      'Softmax', {dim});

you need outputsToSave here, this will make your WASM gradient work

[annxingyuan]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, @Maratyszcza, and @nsthorat)

---

tfjs-backend-wasm/.bazelrc, line 16 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

can you add this to the PR description?

Done

---

tfjs-backend-wasm/src/setup_test.ts, line 40 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

do you have to do this? i think all the kernels needed for the gradient are defined in wasm

Done

---

tfjs-backend-wasm/src/cc/kernels/Softmax.cc, line 85 at r1 (raw file):

Previously, Maratyszcza (Marat Dukhan) wrote…

Why do you hard-code batch size 1?

Done

---

tfjs-backend-wasm/src/cc/kernels/Softmax.cc, line 67 at r2 (raw file):

Previously, Maratyszcza (Marat Dukhan) wrote…

it needlessly increase cache size, but shouldn't affect correctness per se

Done

---

tfjs-backend-wasm/src/cc/kernels/Softmax_test.cc, line 33 at r1 (raw file):

Previously, Maratyszcza (Marat Dukhan) wrote…

out_size > x_size!

Done

---

tfjs-core/src/backends/backend.ts, line 308 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

let's not do this on the backend interface at all, rather we should do this using registerKernel with the string 'Softmax'. Let backend_cpu and webgl have implementations that are modular too.

As discussed, softmax modularization will be part of a follow-up PR: #2751

---

tfjs-core/src/backends/cpu/backend_cpu.ts, line 388 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

no chaining API here and above, use methods directly, e.g. https://github.com/tensorflow/tfjs/blob/master/tfjs-core/src/backends/cpu/backend_cpu.ts#L348

This actually should cause a test failure since we have something that makes sure you don't leave the backend

Done

---

tfjs-core/src/backends/webgl/backend_webgl.ts, line 1600 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

here too make this modular

As discussed, softmax modularization will be part of a follow-up PR: #2751

---

tfjs-core/src/ops/softmax.ts, line 75 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

you need outputsToSave here, this will make your WASM gradient work

Done

[nsthorat]

Reviewed 8 of 8 files at r4.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, @Maratyszcza, and @nsthorat)