Thanks Thomas. I think that may work for my case. I’ve played around with a bit this morning. Just struggling a bit with generalizing it to the 2D case.
I think one of the issues may be, but I am not yet certain, that the src tensor in the scatter operation may only be allowed to be (n,) in shape without any more axes. For example, how would the operation know on what dimension to average over if a tensor of (n, q, r, s) was provided?
I’ll keep looking into it, but just an fyi for others who may have suggestions or input I have not yet solved the problem.