Gradient for tf.gather

[jgartman]

This PR implements the gradient for tf.gather. tensorflow/tfjs#135

---

This change is 

[dsmilkov]

Really nice work! Left a few comments but nothing major. Will LGTM after your reply and merge. Thanks Josh!

---

Reviewed 2 of 2 files at r1.
Review status: all files reviewed at latest revision, all discussions resolved, some commit checks failed.

---

src/ops/array_ops.ts, line 883 at r1 (raw file):

    util.assert(indices.dtype === 'int32', 'Indices must be of dtype `int32`');
    const axes = parseAxisParam(axis, x.shape);

to avoid dealing with axes[0] you can overwrite axis:
axis = parseAxisParam(axis, x.shape)[0]

---

src/ops/array_ops.ts, line 886 at r1 (raw file):

    const grad = (dy: T) => {
      const derX = () => {
        return unsortedSegmentSum(dy, indices, x.shape[axes[0]], axes[0]);

since you already wrote this op and TF has it (tf.unsorted_segment_sum), can you expose it to the public API and add docs?

---

src/ops/array_ops.ts, line 894 at r1 (raw file):

    };
    return ENV.engine.runKernel(
        backend => backend.gather(x, indices, axes[0]), {x, indices}, grad);

to save memory, and since the gradient for indices is meaningless, remove derIndices and remove it from {x, indices} passed to runKernel. That is:

return ENV.engine.runKernel(
        backend => backend.gather(x, indices, axes[0]), {x}, grad);

---

src/ops/array_ops.ts, line 1295 at r1 (raw file):

      segmentIds.dtype === 'int32', 'Segment Ids must be of dtype `int32`');

  const axes = parseAxisParam(axis, x.shape);

same here: axis = parseAxisParam(axis, x.shape)[0]

---

src/ops/array_ops.ts, line 1313 at r1 (raw file):

  const reshapedSegmentIds = ArrayOps.reshape(segmentIds, newShape);
  for (let i = 0; i < numSegments; i++) {
    const f = ArrayOps.fill(x.shape, i, 'int32');

If I understand correctly the shape, you can take advantage of the equal op to broadcast, so you can allocate a single scalar, and also call it segmentId instead of f for readability:
const segmentId = scalar(i, 'int32');

---

Comments from Reviewable

[jgartman]

Review status: 0 of 4 files reviewed at latest revision, 5 unresolved discussions, some commit checks failed.

---

src/ops/array_ops.ts, line 883 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

to avoid dealing with axes[0] you can overwrite axis:
axis = parseAxisParam(axis, x.shape)[0]

Done.

---

src/ops/array_ops.ts, line 886 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

since you already wrote this op and TF has it (tf.unsorted_segment_sum), can you expose it to the public API and add docs?

I exposed unsortedSegmentSum, I wasn't sure exactly where it should live in both the documentation and the ops files. Currently it's under reduction ops in the documentation but in arrays_ops.ts. Let me know if either of those should be changed.

---

src/ops/array_ops.ts, line 894 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

to save memory, and since the gradient for indices is meaningless, remove derIndices and remove it from {x, indices} passed to runKernel. That is:

return ENV.engine.runKernel(
        backend => backend.gather(x, indices, axes[0]), {x}, grad);

Done.

---

src/ops/array_ops.ts, line 1295 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

same here: axis = parseAxisParam(axis, x.shape)[0]

Done.

---

src/ops/array_ops.ts, line 1313 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

If I understand correctly the shape, you can take advantage of the equal op to broadcast, so you can allocate a single scalar, and also call it segmentId instead of f for readability:
const segmentId = scalar(i, 'int32');

Done, thanks for noticing that, I completely missed it.

---

src/ops/array_ops.ts, line 1293 at r2 (raw file):

          CompareOps.equal(segmentId, reshapedSegmentIds).asType('float32');
      const sum = mask.mul(x).sum(axis);
      res.push(sum);

Do I need to dispose of intermediate tensors like mask or do they get cleaned up automatically?

---

Comments from Reviewable

[dsmilkov]

Thanks, and sorry for the delay. Couple small things left (see comments). Really nice work man

---

Reviewed 4 of 4 files at r2.
Review status: all files reviewed at latest revision, 4 unresolved discussions, some commit checks failed.

---

src/ops/array_ops.ts, line 886 at r1 (raw file):

Previously, jgartman (Josh Gartman) wrote…

I exposed unsortedSegmentSum, I wasn't sure exactly where it should live in both the documentation and the ops files. Currently it's under reduction ops in the documentation but in arrays_ops.ts. Let me know if either of those should be changed.

Thanks! The subheading 'Reduction'seems reasonable since it is a reduction operation. As for which file, we havereduction_ops.ts`. Can you move it there?

---

src/ops/array_ops.ts, line 1249 at r2 (raw file):

   * ```js
   * const x = tf.tensor1d([1, 2, 3, 4]);
   * const indices = tf.tensor1d([1, 2, 0, 1]);

this snippet might error unless you explicitly say indices is of type int32

---

src/ops/array_ops.ts, line 1263 at r2 (raw file):

   * calculated. Defaults to 0.
   */
  @doc({heading: 'Operations', subheading: 'Reduction'})

add @operation decorator right above @doc. This will wrap the whole method in a tidy which will dispose of intermediate textures.

---

src/ops/array_ops.ts, line 1293 at r2 (raw file):

Previously, jgartman (Josh Gartman) wrote…

Do I need to dispose of intermediate tensors like mask or do they get cleaned up automatically?

Great q. See my comment above. Add @operation decorator to this method

---

Comments from Reviewable

[jgartman]

Review status: 0 of 6 files reviewed at latest revision, 4 unresolved discussions.

---

src/ops/array_ops.ts, line 886 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Thanks! The subheading 'Reduction'seems reasonable since it is a reduction operation. As for which file, we havereduction_ops.ts`. Can you move it there?

Done.

---

src/ops/array_ops.ts, line 1249 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

this snippet might error unless you explicitly say indices is of type int32

Done.

---

src/ops/array_ops.ts, line 1263 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add @operation decorator right above @doc. This will wrap the whole method in a tidy which will dispose of intermediate textures.

Done.

---

src/ops/array_ops.ts, line 1293 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Great q. See my comment above. Add @operation decorator to this method

Done.

---

Comments from Reviewable

[dsmilkov]

Hey, Josh, sorry for the delay - I was travelling this past week! Thank you!

---

Reviewed 3 of 6 files at r3, 3 of 3 files at r4.
Review status: all files reviewed at latest revision, all discussions resolved, some commit checks broke.

---

Comments from Reviewable