Fuse prelu activation.

[annxingyuan]

Changes

• Add preluActivationWeights to conv and matMul fused kernels

Prelu weights are typically broadcasted from final dimension of x.

E.g. some typical inputs:
x.shape: [1, 64, 64, 16]
a.shape: [1, 1, 16]

Holding off on implementing gradients for fused prelu because the alpha gradient needs access to an intermediate tensor that's buried in the fused kernel.

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

---

This change is 

[nsthorat]

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

---

src/backends/backend.ts, line 134 at r1 (raw file):

  }

  fusedBatchMatMul(

do you expect this to grow anymore? I wonder if it should be a config object now

---

src/backends/webgl/mulmat_packed_gpu.ts, line 26 at r1 (raw file):

  userCode: string;

  // activationWeights

remove this comment

---

src/ops/fused_ops.ts, line 236 at r1 (raw file):

 */
/** @doc {heading: 'Operations', subheading: 'Convolution'} */
function conv2d_<T extends Tensor3D|Tensor4D>(

i feel like we should make this an object now since there are a lot of optional params -- it's probably unlikely people are using these directly so I think we may be able to just update this, what do you think?

---

src/ops/fused_test.ts, line 305 at r1 (raw file):

    const result = tf.fused.conv2d(
        x, w, stride, pad, 'NHWC', [1, 1], null, null, 'prelu', alpha);

name the nulls here and elsewhere

[annxingyuan]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @dsmilkov and @pyu10055)

---

src/backends/backend.ts, line 134 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

do you expect this to grow anymore? I wonder if it should be a config object now

Done

---

src/backends/webgl/mulmat_packed_gpu.ts, line 26 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

remove this comment

Done

---

src/ops/fused_ops.ts, line 236 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

i feel like we should make this an object now since there are a lot of optional params -- it's probably unlikely people are using these directly so I think we may be able to just update this, what do you think?

Done

---

src/ops/fused_test.ts, line 305 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

name the nulls here and elsewhere

Done - no longer needed after creating config object

[dsmilkov]

Nice work! Just a couple comments about documenting.

Reviewed 7 of 7 files at r2.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, and @pyu10055)

---

src/ops/fused_ops.ts, line 42 at r2 (raw file):

 * ```
 *
 * @param obj Configuration

same here --- see the other comment below.

---

src/ops/fused_ops.ts, line 224 at r2 (raw file):

 * bias and applying an activation.
 *
 * @param obj Configuration

Say "An object with the following properties" and place a dash (-) in front of each property so markdown renders it nicely as a list.

---

src/ops/fused_ops.ts, line 254 at r2 (raw file):

 *   bias Tensor to be added to the result.
 *   activation Name of activation kernel (defaults to `linear`).
 *   preluActivationWeights Tensor of prelu weights.

say a bit more. e.g. what shape are these weights, what does this do (prelu is applied after the output of the conv) etc.

---

src/ops/fused_util.ts, line 29 at r2 (raw file):

  bias?: Tensor,
  activation?: Activation,
  preluActivationWeights?: Tensor

add a comment for each field with short description. This would be especially useful for preluActivationWeights since many readers will not know what that is.

[dsmilkov]

Actually, one comment about API that we can chat offline.

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, and @pyu10055)

---

src/ops/fused_ops.ts, line 257 at r2 (raw file):

 */
/** @doc {heading: 'Operations', subheading: 'Convolution'} */
function conv2d_<T extends Tensor3D|Tensor4D>({

i'm worried that our API is departing from tf.nn.conv2d. The way we resolved this with matmul is we have a separate tf.fused.matmul. cc @nsthorat I think we should have tf.fused.conv2d for consistency.

[nsthorat]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, and @pyu10055)

---

src/ops/fused_ops.ts, line 257 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

i'm worried that our API is departing from tf.nn.conv2d. The way we resolved this with matmul is we have a separate tf.fused.matmul. cc @nsthorat I think we should have tf.fused.conv2d for consistency.

this is already under the fused namespace

[dsmilkov]

Reviewable status: complete! 2 of 1 approvals obtained (waiting on @annxingyuan and @pyu10055)

---

src/ops/fused_ops.ts, line 257 at r2 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

this is already under the fused namespace

Ah, i missed the filename.

[annxingyuan]

Reviewable status: complete! 2 of 1 approvals obtained (waiting on @annxingyuan and @pyu10055)

---

src/ops/fused_ops.ts, line 42 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

same here --- see the other comment below.

Done

---

src/ops/fused_ops.ts, line 224 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Say "An object with the following properties" and place a dash (-) in front of each property so markdown renders it nicely as a list.

Done

---

src/ops/fused_ops.ts, line 254 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

say a bit more. e.g. what shape are these weights, what does this do (prelu is applied after the output of the conv) etc.

Done

---

src/ops/fused_ops.ts, line 257 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Ah, i missed the filename.

Done