What's Changed
* Bump jax from 0.4.31 to 0.4.32 in /mpi4jax/_src by dependabot in https://github.com/mpi4jax/mpi4jax/pull/255
* Update _latest_jax_version.txt to jax 0.4.33 by PhilipVinc in https://github.com/mpi4jax/mpi4jax/pull/258
* Fix docs by dionhaefner in https://github.com/mpi4jax/mpi4jax/pull/259
* Ensure out_types is an interable on gpu by dionhaefner in https://github.com/mpi4jax/mpi4jax/pull/261
**Full Changelog**: https://github.com/mpi4jax/mpi4jax/compare/v0.5.3...v0.5.4