JAX for Scientific Computing: Bridging the gap between performance, productivity and portability

In this talk, I will discuss the use of the JAX library for solving problems in scientific computing. After a brief discussion of the main features of the library, I will focus on how these features enable researchers to develop more efficient implementations of their scientific models that can run on both CPUs and GPUs and that can interface with machine-learning models seamlessly. I will show examples of open-source JAX-based solvers for a few problems across computational science and engineering and share personal experience of developing software for applications in fluid mechanics and plasma physics. Finally, I will discuss some limitations of the JAX library and issues and shortcomings that I have experienced.Acknowledgements I am grateful for support from the Fusion Computing Lab, a collaboration between UKAEA and STFC Hartree Centre.A recording of this session is available on YouTube: https://youtu.be/NxuREjRpEzk