I've seen it although haven't dived deep into it. It looks like they have some interesting support for particle cell data structures, but is fairly complicated and carries limitations: https://jax-md.readthedocs.io/en/main/_modules/jax_md/partit...

Last time I looked at JAX MD it didn't support most of the force field terms necessary for simulating proteins and DNA. For example, it could do n-body simulations with some potential, but not the bonds/torsions between atoms. It's unclear if they added support, but that's a huge gap in functionality compared to other systems.