Common utilities for working with jax-md.

In order to use these in a colab notebook run the following line at the top of your notebook:

!pip install -q git+https://github.com/cpgoodri/goodgrp_utils.git