Latest versions of JAX and objax cause compile slow down

[wil-j-wil]

It is recommended to use the following versions of jax and objax:

jax==0.2.9
jaxlib==0.1.60
objax==1.3.1

This is because of this objax issue which causes the model to JIT compile "twice", i.e. on the first two iterations rather than just the first. This causes a bit of a slow down for large models, but is not an problem otherwise.