Example notebook with PMAP sampling of RBMs trained on MNIST digits

[StannisZhou]

Resolves the first todo in #77. Also includes a simple change with which we can index 1D variable arrays with integers (instead of length 1 tuples).

Currently a lot of time is spent on graph construction (adding factors and constructing factor wirings) since PGMax is not making use of any of the structures present in an RBM.

In the future, to improve usability, it would be helpful to:

1. See if we can find an efficient way to parallelize many of the involved for loops (Optimize and parallelize structure compiling function #18)
2. Introduce a mechanism to serialize a model to disk
3. Improve interface so that we can explicitly use already constructed wiring.

Separately, it would be helpful to introduce a pure functional interface so that we can leverage functionalities like vmap in JAX (which would allow us to parallelize sampling multiple digits in this use case). Connected to #68

See below for some example samples


:

[NishanthJKumar]

Looks mostly good, but 1 potentially-important comment about dependencies and another more minor comment about testing. Will approve after these are addressed/replied to!

[NishanthJKumar]

LGTM! Thanks for addressing the previous comments so quickly 😄

[StannisZhou]

Looks like the use of __file__ would break the notebooks. Having the notebooks running seems more important so I got rid of the __file__ and also excluded heretic model example from unit tests for now.