JAX MD: A Framework for Differentiable Physics - Schoenholz, Samuel; Cubuk, Ekin Dogus (Advances in Neural Information Processing Systems, 2020)