Expose GPGPUProgram but do it right this time

[aman-tiwari]

Description

See #1202

---

For repository owners only:

Please remember to apply all applicable tags to your pull request.
Tags: FEATURE, BREAKING, BUG, PERF, DEV, DOC, SECURITY

For more info see: https://github.com/tensorflow/tfjs/blob/master/DEVELOPMENT.md

---

This change is 

[dsmilkov]

Reviewable status: 0 of 1 approvals obtained (waiting on @aman-tiwari)

---

src/kernels/webgl/webgl_custom_op_test.ts, line 64 at r2 (raw file):

    const backward = (dy: T) => {
      return {
        x: () => tf.tidy(() => {

You don't need tidy since we should wrap that for you in the engine when we backpropagate.

---

src/kernels/webgl/webgl_custom_op_test.ts, line 71 at r2 (raw file):

    };

    const res = tf.ENV.engine.runKernel(forward, {x}, backward);

Instead of runKernel, how about using tf.customGrad() which is already documented and aligned with https://www.tensorflow.org/api_docs/python/tf/custom_gradient

[aman-tiwari]

Reviewable status: 0 of 1 approvals obtained (waiting on @aman-tiwari)

---

src/kernels/webgl/webgl_custom_op_test.ts, line 71 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Instead of runKernel, how about using tf.customGrad() which is already documented and aligned with https://www.tensorflow.org/api_docs/python/tf/custom_gradient

customGrad only works if the things inside it call runKernel at some point (according to @nsthorat)

[dsmilkov]

Reviewable status: 0 of 1 approvals obtained (waiting on @aman-tiwari and @dsmilkov)

---

src/kernels/webgl/webgl_custom_op_test.ts, line 64 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

You don't need tidy since we should wrap that for you in the engine when we backpropagate.

Move the construction of the backpropProgram outside of the derX function so we don't recreate it every time we backprop.

---

src/kernels/webgl/webgl_custom_op_test.ts, line 71 at r2 (raw file):

Previously, aman-tiwari (Aman Tiwari) wrote…

customGrad only works if the things inside it call runKernel at some point (according to @nsthorat)

it should work without runKernel, unless there is a bug. Anyhow, runKernel() is good for now, but we should revisit customGrad() and figure out what's going on.

[aman-tiwari]

Reviewable status: 0 of 1 approvals obtained (waiting on @aman-tiwari)

---

src/kernels/webgl/webgl_custom_op_test.ts, line 71 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

it should work without runKernel, unless there is a bug. Anyhow, runKernel() is good for now, but we should revisit customGrad() and figure out what's going on.

@nsthorat is looking at the bug now

[dsmilkov]

Reviewed 1 of 1 files at r3.
Reviewable status: 0 of 1 approvals obtained (waiting on @dsmilkov)