Implement gradient operator for GatherByKeys. (#24348)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/24348
Partition + GatherByKeys pair is pretty handy for implementing strategy where
part of the keys will be on local machine, while part of the keys will end up
on the remote machin (for cases when there is exactly 1 id).
Reviewed By: aazzolini
Differential Revision: D16802988
fbshipit-source-id: 4c7ac97fc0db3ce88575fccab0c7bf69dcbef965
