jit model init
Created by: PhilipVinc
This should speed-up the initialisation (Creation of variational state) for very large models, because under jit we won't be evaluating the output (instead without jit it is evaluated and then discarded). It will add a tiny slowdown to initialisation for smaller models.