Bump jax from 0.6.2 to 0.7.0 in /mpi4jax/_src (#283)

dependabot[bot] · dionhaefner · web-flow · commit ba8e3a7985f9 · 2025-07-27T20:54:45.000+02:00
* Bump jax from 0.6.2 to 0.7.0 in /mpi4jax/_src Bumps [jax](https://github.com/jax-ml/jax) from 0.6.2 to 0.7.0. - [Release notes](https://github.com/jax-ml/jax/releases) - [Changelog](https://github.com/jax-ml/jax/blob/main/CHANGELOG.md) - [Commits](jax-ml/jax@jax-v0.6.2...jax-v0.7.0) --- updated-dependencies: - dependency-name: jax dependency-version: 0.7.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> * instantiate zeros * bump python version for coverage test --------- Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Dion Häfner <mail@dionhaefner.de>
diff --git a/.github/workflows/covecov-coverage.yml b/.github/workflows/covecov-coverage.yml
@@ -15,8 +15,9 @@ jobs:
       fail-fast: false
 
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.13"]
         os: [ubuntu-latest]
+        mpi: [openmpi]
 
     env:
        MPICH_INTERFACE_HOSTNAME: localhost
diff --git a/mpi4jax/_src/_latest_jax_version.txt b/mpi4jax/_src/_latest_jax_version.txt
@@ -1,3 +1,3 @@
 # latest tested JAX version
 # bumped automatically by dependabot
-jax==0.6.2
+jax==0.7.0
diff --git a/mpi4jax/_src/collective_ops/allreduce.py b/mpi4jax/_src/collective_ops/allreduce.py
@@ -223,6 +223,7 @@ def mpi_allreduce_transpose_rule(x_tan, *x_args, op, comm, transpose):
         raise NotImplementedError(
             "The linear transpose of allreduce is only defined for op=MPI.SUM"
         )
+    x_tan = ad.instantiate_zeros(x_tan)
     res = mpi_allreduce_p.bind(x_tan, op=op, comm=comm, transpose=(not transpose))
     return (res,)
 

Original file line number	Diff line number	Diff line change
`@@ -223,6 +223,7 @@ def mpi_allreduce_transpose_rule(x_tan, *x_args, op, comm, transpose):`
`223`	`223`	`raise NotImplementedError(`
`224`	`224`	`"The linear transpose of allreduce is only defined for op=MPI.SUM"`
`225`	`225`	`)`
	`226`	`+ x_tan = ad.instantiate_zeros(x_tan)`
`226`	`227`	`res = mpi_allreduce_p.bind(x_tan, op=op, comm=comm, transpose=(not transpose))`
`227`	`228`	`return (res,)`
`228`	`229`