`apply` method really slow on CPU about dm-haiku

[tomhennigan]

Hey @Guillem96, I just tested this in Google Colabatory (CPU only) and your init function takes a few seconds initially (I believe this is JAX JIT compiling kernels in the background for things like the random initial parameter values etc). If I run it a second time it only takes 778ms:

>>> %time params = vgg_forward_fn.init(rng, np.expand_dims(im, 0))
CPU times: user 9.37 s, sys: 366 ms, total: 9.74 s
Wall time: 8.78 s

>>> %time params = vgg_forward_fn.init(rng, np.expand_dims(im, 0))
CPU times: user 1.09 s, sys: 158 ms, total: 1.25 s
Wall time: 778 ms

If we JIT it then it runs in 10ms:

>>> jit_init = jax.jit(vgg_forward_fn.init)
>>> jit_init(rng, np.expand_dims(im, 0))  # Trigger JIT compilation for our input
>>> %time params = jit_init(rng, np.expand_dims(im, 0))
CPU times: user 9.85 ms, sys: 36 µs, total: 9.88 ms
Wall time: 9.31 ms

So I'm not sure what is going on here, but suspect it is perhaps related to your JAX version or something with JAX + WSL? As far as I can tell Haiku is functioning correctly.