GitHub repository: AlanPearl/multigrad
Primary language: Python
Description: Framework to implement JAX models that can be distributed over MPI