stop gradient on allreduce create_token (#30)

* stop gradient on allreduce create_token
mpi4jax · Sep 23, 2020 · d8f8434 · d8f8434
1 parent 10551c1
commit d8f8434
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 45 deletions.
diff --git a/mpi4jax/__init__.py b/mpi4jax/__init__.py
@@ -12,7 +12,6 @@
 from .collective_ops.send import Send
 from .collective_ops.sendrecv import Sendrecv
 from .flush import flush
-from .token import create_token
 from .warn import disable_omnistaging_warning
 
 # at exit, we wait for all pending operations to finish
@@ -27,7 +26,6 @@
     "Send",
     "Recv",
     "Sendrecv",
-    "create_token",
     "disable_omnistaging_warning",
     "flush",
 ]

diff --git a/mpi4jax/collective_ops/allreduce.py b/mpi4jax/collective_ops/allreduce.py
@@ -6,7 +6,7 @@
 from jax.interpreters import ad, xla
 from jax.lib import xla_client
 
-from ..token import create_token
+from jax.lax import create_token, stop_gradient
 from ..utils import (
     HashableMPIType,
     _constant_s32_scalar,
@@ -52,7 +52,7 @@ def Allreduce(x, op, comm=_MPI.COMM_WORLD, token=None):
             This result can be ignored if result forces a data dependency.
     """
     if token is None:
-        token = create_token(x)
+        token = create_token(stop_gradient(x))
 
     op = wrap_as_hashable(op)
     comm = wrap_as_hashable(comm)

diff --git a/mpi4jax/token.py b/mpi4jax/token.py