Jax-md

This release contains a large refactor of the JAX MD simulation environments as well as new features built on this refactor including support for rigid body physics and preliminary support for simulations on TPU. This version also contains a preliminary version of Nequip as well as code to deserialize saved models.

- Refactor of simulation environments and rigid body physics with (ellamichelleking, xiyudu, mc2engel, rkruegs123)
- Preliminary TPU code with Sohl-Dickstein.
- Neural network infrastructure + nequip with simonbatzner, doguscubuk, amilmerchant.
- Improved handling of neighbor lists in triclinic systems.

* Adds an equivariant graph neural network example.
* Miscellaneous fixes to the renderer (added a non-verbose mode and fixed the render frame inference in the case of a static snapshot).
* Neighbor lists now default to `dr_threshold = 0`.
* Add `ParameterTree` construct to `smap` functions. These allow users to pass PyTrees of data into `smap` functions with control over how those parameters are mapped among particles.

- Small fixes to neighbor list / cell list capacity logic.
- Small fix to NPT simulations.
- Tersoff model added by SangHyukYoo.
- Documentation improvements (with more contributions from adrhill).

NPT Fixes, documentation, and neighbor list improvements.

Changes:

- Significant documentation improvements by adrhill.
- Add custom mask function to neighbor lists by oliverdutton.
- Small fixes to renderer.
- Fixes to NPT simulation environment using fixed equation for pressure in v0.1.23.
- Add SyMD (symmetric molecular dynamics) example notebook based on https://github.com/whitead/symd.

Improves and adds support for a number of mechanical properties.

Changes:

- Adds `perturbation` parameter to all spaces.
- Fixes pressure computation.
- Adds stress tensor computation to `quantity.py`.
- Carl Goodrich (cpgoodri) added support for athermal linear elasticity computations in `elasticity.py`, along with a great colab notebook!
- Patrick Marks (pmarks) improved neighbor lists by switching to int32 for indices.
- Maxi Lechner (MaxiLechner) added an implicit differentiation notebook combining JAX MD with JaxOpt.

More neighbor list improvements.

Changes:

- Improves overflow checking in neighbor lists by also checking for buffer overflows in cell lists.
- Modernizes cell list API to mimic neighbor lists (adds `allocate` and `update` functions).
- Neighbor list no longer carries the cell list update function explicitly.
- Thread neighbor list keyword arguments through energy functions.
- Remove scaling of `dr_threshold` by radius in soft sphere and Lennard-Jones potentials.