A large fraction of computational science involves simulating the dynamics of particles that interact via pairwise or many-body interactions. These simulations, called Molecular Dynamics (MD), span a vast range of subjects from physics and materials science to biochemistry and drug discovery. Most MD software involves significant use of handwritten derivatives and code reuse across C++, FORTRAN, and CUDA. This is reminiscent of the state of machine learning before automatic differentiation became popular. In this work we bring the substantial advances in software that have taken place in machine learning to MD with JAX, M.D. (JAX MD). JAX MD is an end-to-end differentiable MD package written entirely in Python that can be just-in-time compiled to CPU, GPU, or TPU. JAX MD allows researchers to iterate extremely quickly and lets researchers easily incorporate machine learning models into their workflows. Finally, since all of the simulation code is written in Python, researchers can have unprecedented flexibility in setting up experiments without having to edit any low-level C++ or CUDA code. In addition to making existing workloads easier, JAX MD allows researchers to take derivatives through whole-simulations as well as seamlessly incorporate neural networks into simulations. This paper explores the architecture of JAX MD and its capabilities through several vignettes. Code is available at www.github.com/jaxmd/jax-md along with an interactive Colab notebook that goes through all of the experiments discussed in the paper.
This paper announces a new software package for simulating molecular dynamics which includes close integration with a neural network / machine learning library--the first to do so.  Straightforward access to hardware acceleration (e.g. GPU) is provided for both the simulation and machine learning.
I lean toward accepting this submission.  If it were only about simulation molecular dynamics using hardware accelerators, I would question the appropriateness of the venue, but because it is explicitly intended to support training and usage of learned potential functions, it seems suitable.  Still might be better placed in a physics/chemistry venue, as where most of the references come from and likely where users would, too. The application area is no doubt an important research technique.  The paper is clearly written, with enough specific examples to contrast previous pain points in this line of work against its smoother interface.  
All of these points are fine for a package-release/tutorial paper, but for a conference paper, might hope to see these addressed:
Description of the elements of the design of JAX which are useful here are presented, and appear distinct from other AD libraries like Tensorflow or PyTorch, although the authors stop short of explicitly stating which functionality would be more difficult/impossible to support with the possible alternatives (automatic vectorization of the simulations seems like one?).
Limitations of the library and drawbacks of any design decisions (something must be traded at some point?) are not explicitly mentioned.  
Despite mentioning numerous existing MD libraries, no performance comparison is drawn against any other.  
Could also show some demonstration of running an experiment which has complexity on par with state of the art research?  The bubble raft example is great for illustrative purposes, but it could be better to save some of that for a tutorial and use space to exercise this library on a relevant problem and show performance there.
I took a quick glance at the code on github; it is substantial but not huge, cleanly organized, and is well-documented.