What's Changed * Bump jax from 0.3.25 to 0.4.0 in /mpi4jax/_src by dependabot in https://github.com/mpi4jax/mpi4jax/pull/182 * Bump jax from 0.4.0 to 0.4.1 in /mpi4jax/_src by dependabot in https://github.com/mpi4jax/mpi4jax/pull/183
What's Changed * Port from XLA CustomCall to MLIR by dionhaefner in https://github.com/mpi4jax/mpi4jax/pull/179 * Fixed a bug that led to incorrect results whenever XLA eliminated consecutive transposes (176). * Now requires JAX >= 0.3.25.