Hi, can you specify what you mean by “keeping the gradient of the original tensor”?
If it means that you just want to propagate the gradient incoming to:
[[[A], 
  [C]]]

we can note:
[[[grad(A)], 
  [grad(C)]]]

to the input of the operation in the most straight forward way you would obtain:
[[[grad(A)], 
  [None],
  [grad(C)]]]

Which is what would torch.gather do.
In case you want something else, please specify but I assume you will have to extend torch.autograd as in:
https://pytorch.org/docs/stable/notes/extending.html
Best,
Samuel