I am interested in implementing Exponential Moving Average that would allow running backward() on it, in such way that it could be applied to tensors with substantial graphs creating them.
The straightforward implementations create an expanding graph that includes all graphs that create the past versions of the averaged tensor, and running backward() quickly runs out of memory.
I understand that this has already been implemented in batch normalization where the means and variances of the inputs to the layer are kept as exponential moving averages.  However, it is written in C++ and I could not dig out the implementation in the code.
I think the idea is to keep all the gradients of the moving average tensor w.r. to all the leaves together with the moving average (rather than with the leaf tensors); this would allow easy calculation of the gradients together with the moving average update.
However, the way of implementing it in python that I see, is tricky and inefficient, by subclassing Tensor to keep all the leaf gradients with it, and overriding Tensor.backward().
Any better idea?