Accumulation of Monte-Carlo gradients within a flax module to avoid OOM error

[etienne-thuillier]

I noticed that I posted this at the wrong place. And so I copied the question in flax's Q&A
google/flax#2301

I did not find a way to delete this thread...

[YouJiacheng]

Check this for how to convert a flax.module to pure functions init and apply:
https://flax.readthedocs.io/en/latest/design_notes/lift.html#functionalization

class Model(nn.Module):
  @nn.compact
  def __call__(self, ...):
    ...
    decoder = Decoder(sigma_floor=1.0e-3, parent=None)
    apply_fn = decoder.apply
    decoder_params = self.param('mlp', decoder_params, decoder.init)
    if self.is_mutable_collection('params'): # Optional, avoid useless apply_fn in initialization
        apply_fn = lambda: _, x: dummy_eval_of_decoder(x)
        # dummy eval only need to output same shape and dtype
        # dummy_eval_of_decoder(z).shape == apply_fn(decoder_params, z).shape
        # dummy_eval_of_decoder(z).dtype == apply_fn(decoder_params, z).dtype
    # now apply_fn is a pure function full compatible with JAX

You can also use lifted transformation provided by Flax
https://flax.readthedocs.io/en/latest/_autosummary/flax.linen.custom_vjp.html

[etienne-thuillier]

Sorry @YouJiacheng, I noticed your answer only after I copied my post to flax's Q&A :|

I am going to put a link to this current original post so that your answer can be consulted.

Thank you for the tip. I'll have to do some reading before I can understand it!