Workshop: From Molecules to Materials: ICLR 2023 Workshop on Machine learning for materials (ML4Materials)
JAX-XC: Exchange Correlation Functionals Library in Jax
Kunhao Zheng · Min Lin
We present JAX-XC, an open-source library that provides exchange-correlation functionals in Jax. JAX-XC is built from LIBXC, its correctness has been verified numerically against LIBXC. Thanks to Jax, JAX-XC is end-to-end differentiable, computationally more efficient thanks to the vectorization provided by XLA, and also portable on various accelerators. More importantly, as more research is focusing on machine learning for density functional theory, we hope that JAX-XC could serve as a deep learning-friendly tool and a stepping-stone for researchers working in the intersection of deep learning and density functional theory.