Jit the loss or the loop?

[mathDR]

Would it be better to Jit a lax.scan of the entire optimization loop versus jitting the loss and using a regular python loop?

I see that usually the latter is done in most examples

[AlexeyKurakin]

We usually do Jit of individual training steps to keep code clean and easy to understand.
Also keep in mind that if you Jit multiple iterations of training loop you will have to pre-fetch several batches from the dataset beforehand.

Overall, if you Jit multiple training steps then you probably will get some performance boost. I think people don't do this because for any complex model and non-trivial dataset this performance boost would be relatively small and at the same time code will become more complex.

[mathDR]

Oh I see your point about fetching data. This is true for Neural network models. My use case is for optimizing hyperparameters of a Gaussian process though. I would think for models that do not need prefetched data, using a JAX fori_loop and JIT compiling it would see a pretty large performance boost, right?

In any case, I will run an MVP and test the timings!