I talked to the team, but unfortunately, jax-md at the ttime didn't do bond angles or torsions, so it wasn't good for biomolecular simulations.
My work mostly predated tensorflow and was much more about massive-scale embarassingly parallel computing, and produced some interesting large-scale results from MD and protein folding.
My work mostly predated tensorflow and was much more about massive-scale embarassingly parallel computing, and produced some interesting large-scale results from MD and protein folding.
https://www.nature.com/articles/nchem.1821 https://onlinelibrary.wiley.com/doi/full/10.1002/pro.2389