mpi4jax: Zero-copy MPI communication of JAX arrays
Dion Häfner, Filippo Vicentini · The Journal of Open Source Software · 2021
The tensor framework JAX (Bradbury et al., 2018) combines expressivity and performance while retaining an accessible pure Python interface.Expressivity is achieved by treating functions as first-class objects, while efficiency is obtained by compiling to machine code justahead-of-time.* Contributed equally, order determined by coin flip.