We have just released a Jax implementation of stochastic OnsagerNet, with applications to polymer dynamics.

https://github.com/MLDS-NUS/onsagernet-jax

References

[1] X. Chen et al., “Constructing custom thermodynamics using deep learning,” Nature Computational Science, vol. 4, no. 1, pp. 66–85, 2024. Available: https://www.nature.com/articles/s43588-023-00581-5

[2] K. S. Novoselov and Q. Li, “Learning physical laws from observations of complex dynamics,” Nat Comput Sci, pp. 1–2, Jan. 2024, doi: 10.1038/s43588-023-00590-4.