some (very) small suggestions on the examples

[murphyk]

• In examples/gmrf.ipynb, maybe replace from jax.example_libraries import optimizers with optax. Also the term "GMRF" usually refers to Gaussian MRF, not grid :) Finally the learning code is super slow, even on an A100... (public colab)
• in examples/rbm.ipynb, you mix jax.vmap with np.random, which seems dangerous...
• in examples/ising_model.ipynb, the line batch_loss = jax.jit(jax.vmap(loss, in_axes=(None, {variables: 0}), out_axes=0)) is very obscure. What is the dict with a key called 'variables'? And in what sense are these evidence / potential 'updates', as opposed to just static things?
• it would be helpful to add a link to the ising example from https://github.com/deepmind/PGMax/tree/main#getting-started

[antoine-dedieu]

Thank you @murphyk for these detailed comments!

1. We have renamed the gmrf.py notebook to grid_mrf.py, to avoid confusion. Instead of training a GridMRF model from scratch, which was really slow, we now simply perturb the weights of the pretrained model. We only finetune the weights for 500 samples, starting from these perturbed weights. On a Colab GPU instance with 16 Gb of RAM, finetuning should now take around 5 minutes and should only require 11 Gb of RAM.

2. We are now properly using jax.random with jax.vmap in all our notebooks

3. We have removed these obscure lines.

4. We have modified the Ising example and are now introducing it in our README as a tutorial, so that users can easily get started with PGMax.