Fixed the iMAML tutorial

[zaccharieramzi]

Basically implemented the correct iMAML training setup:

• corrected phase in sinus
• made sure to use more than 4 tasks for meta-training
• used vmap for outer loss computation

Then adapted all the plots.

TPU testing remains to be done.

This fixes #417 .

@fabianp don't hesitate to tell me already if some things should be improved w.r.t for example using LBFGS as the inner solver, a more complex MLP (with swish activation), ...

[zaccharieramzi]

Unfortunately I cannot test on TPU on Colab because they are not supported anymore:

RuntimeError: 
As of JAX 0.4.0, JAX only supports TPU VMs, not the older Colab TPUs.

We recommend trying Kaggle Notebooks
(https://www.kaggle.com/code, click on "New Notebook" near the top) which offer
TPU VMs. You have to create an account, log in, and verify your account to get
accelerator support.
Once you do that, there's a new "TPU 1VM v3-8" accelerator option. This gives
you a TPU notebook environment similar to Colab, but using the newer TPU VM
architecture. This should be a less buggy, more performant, and overall better
experience than the older TPU node architecture.

It is also possible to use Colab together with a self-hosted Jupyter kernel
running on a Cloud TPU VM. See
https://research.google.com/colaboratory/local-runtimes.html
for details.

@fabianp Would you like me to test on Kaggle even if the quick button still points to Colab?
Otherwise, tested on GPU it is twice as fast, so working I'd say.

[zaccharieramzi]

If anyone is interested, I live streamed the last part of my implementation of this PR on twitch: https://www.twitch.tv/videos/1807716437

[mblondel]

Thanks a lot for the contribution @zaccharieramzi.

@fabianp Can you take a look when you get some time? Thanks!

[fabianp]

sure will do. the streaming is awesome 😄

[fabianp]

thanks a lot for this @zaccharieramzi . Could you please run jupytext as described here (https://jaxopt.github.io/stable/developer.html#syncing-notebooks) so that we can see more easily the diff ?

[zaccharieramzi]

@fabianp sure done!
I also added jupytext to the docs requirements as it was not present.
An alternative to syncing the 2 versions each time would be to use ReviewNB (I use it in some of my repos it works really well): https://www.reviewnb.com/. Don't know how easy it would be to put in place though with the 2 versioning systems.

[fabianp]

Thanks a lot @zaccharieramzi, this looks great! And thanks for the link to reviewnb, it does look great indeed. We use jupytext mostly because JAX uses this same tooling.

One last thing before we merge. Could you please squash these into a single commit ? (it makes it easier for us to sync with the Google infra)

[zaccharieramzi]

@fabianp sure done!