  Electronic structure simulation (ESS) has been used for decades to provide
quantitative scientific insights on an atomistic scale, enabling advances in
chemistry, biology, and materials science, among other disciplines. Following
standard practice in scientific computing, the software packages driving these
studies have been implemented in compiled languages such as FORTRAN and C.
However, the recent introduction of machine learning (ML) into these domains
has meant that ML models must be coded in these languages, or that complex
software bridges have to be built between ML models in Python and these large
compiled software systems. This is in contrast with recent progress in modern
ML frameworks which aim to optimise both ease of use and high performance by
harnessing hardware acceleration of tensor programs defined in Python. We
introduce MESS: a modern electronic structure simulation package implemented in
JAX; porting the ESS code to the ML world. We outline the costs and benefits of
following the software development practices used in ML for this important
scientific workload. MESS shows significant speedups n widely available
hardware accelerators and simultaneously opens a clear pathway towards
combining ESS with ML. MESS is available at
https://github.com/graphcore-research/mess.
