Refactor utils.mpi
Created by: PhilipVinc
This is just the refactoring of utils.mpi
part of #690
I think we can merge only this one for now. It gives us a place to put all mpi-related functions, gives different names to mpi4py and mpi4jax functions so that we avoid bugs in the future. It also uses this to fix the driver's callback bug