Jit'ing the grad function

[srush]

I have an inner function in my network where calling jit(grad(fn)) seems to yield a big memory savings. However in the standard flax pattern of training I don't really see how to do this. Since we call .apply each time, I would need to run jit in the inner loop, which seems bad.

Is there a way to do this that I'm missing?

[marcvanzee]

(I am not sure if I fully understand your question, so please clarify if my answer below is not what you were asking.)

Do you mean you are calling grad inside your Module's apply function? If you jit something inside another jit block, then the inner jit should be a no-op. We usually jit the entire train function, which calls the grad function so that will be jitted inside this bigger block (we actually prefer jitting bigger blocks because it gives XLA more opportunity for optimizing things -- at the cost of longer compile time).

Since we call .apply each time, I would need to run jit in the inner loop, which seems bad.

The jitted function will be compiled only once for each shape it is called with, so if the shapes are the same in the loop then this is not bad.

[srush]

Thanks! I missed that the whole train step was jit'd.