[tfjs-core] fix gather gradient when batchDims is 1

[paradite]

This PR fixes #7494, where error is encountered when computing gradient of tf.gather with batchDims=1.

I am opening a PR with a naive and unoptimized fix first to see if the CI passes, and gather feedback on the direction of the fix.

If the CI passes, I will work on optimizing the logic if needed.

I am not really familiar with the maths behind the original implementation in derX of gatherGradConfig, and I couldn't get it to work within the derX function, I am open to suggestions on better and more optimized way to fix the issue.

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

[mattsoulanille]

/gcbrun

[mattsoulanille]

Thanks for the contribution! It looks like you're splitting the batched tensor, applying the gradient function to each batch, and then stacking them together again. This seems reasonable to me.

The original math for derX is also not really clear to me. For y = tf.gather(x, indices, axis, batchDim), I'd expect dy/dx to just copy the selected values specified by the indices variable from the dy variable and be zero everywhere else. However, it seems to be using unsortedSegmentSum. I think this is to account for repeated indices. There's also some complicated transpose logic to make the input work for unsortedSegmentSum. It might be possible to get the indices of these two ops to work correctly for arbitrary batch dimensions, but I'm not sure. I'm fine with your approach, and we can revisit it if we need better performance in the future.

There's probably a way to apply your approach for generic batchDims, but I'm not going to block this PR because of that.

Also, it looks like CI is failing due to code linting. I'll push a patch to fix that.

[mattsoulanille]

/gcbrun

[mattsoulanille]

LGTM

[paradite]

Thank you @mattsoulanille for the help and review. I will take note of the lint issue for future PRs.

Hi @pyu10055, can I get your blessing and help to merge this?