I want to keep a state between the updates (imagine simple L2, but with different multiplier on each iteration). Which is think is challenging to do in a clean way.
Could you please give me an example, say, basic L2 (or at least something)?
This factor primitive is not particularly documented (neither in numpyro nor in pyro). Best i could find is this Custom Loss Function Implementation - #4 by fritzo and the author uses it inside the guide, not the objective per se.