Moving to a functional interface

[StannisZhou]

This PR moves PGMax to a functional interface.

Main changes to the interface (resolves #82):

• Instead of generating init_msgs from FactorGraph and rely on a member function (run_bp) to run belief propagation, we derive an immutable FactorGraphState based on added factors to a factor graph, and use the FactorGraphState to initialize a BPState.
• A BP function then takes a BPState and generates 3 pure functions:
  1. run_bp, which takes log_potentials/ftov_msgs/evidence updates specified in regular python dictionaries, and outputs a BPArrays containing results after BP.
  2. get_bp_state which constructs a BPState from a BPArrays
  3. get_beliefs which calculates beliefs from a BPArrays.

The functional approach allows easy use of regular function transformations in JAX like vmap, grad. See examples/ising_model.py and examples/rbm.py for some added examples on using vmap and grad with the new functional interface. Resolves #68.

Also made various refactors/cleanups (e.g. for unit tests/docstrings/naming, etc.). Resolves #75. Resolves #71.

[NishanthJKumar]

LGTM after last couple minor comments addressed! Great work - really exciting to see the codebase this mature!