mpi4jax
diff --git a/‎.github/workflows/mpi-tests.yml
Lines changed: 1 addition & 1 deletion b/‎.github/workflows/mpi-tests.yml
Lines changed: 1 addition & 1 deletion
diff --git a/‎.github/workflows/pre-commit.yml
Lines changed: 3 additions & 3 deletions b/‎.github/workflows/pre-commit.yml
Lines changed: 3 additions & 3 deletions
diff --git a/‎docs/sharp-bits.rst
Lines changed: 0 additions & 38 deletions b/‎docs/sharp-bits.rst
Lines changed: 0 additions & 38 deletions
diff --git a/‎docs/usage.rst
Lines changed: 5 additions & 20 deletions b/‎docs/usage.rst
Lines changed: 5 additions & 20 deletions
diff --git a/‎examples/shallow_water.py
Lines changed: 20 additions & 29 deletions b/‎examples/shallow_water.py
Lines changed: 20 additions & 29 deletions
@@ -30,7 +30,7 @@ jobs:
           - os: ubuntu-latest
             python-version: "3.10"
             mpi: openmpi
-            jax-version: "0.4.35"
+            jax-version: "0.5.1"
 
     env:
        MPICH_INTERFACE_HOSTNAME: localhost
 
@@ -12,6 +12,6 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
-    - uses: actions/setup-python@v2
-    - uses: pre-commit/action@v2.0.0
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
+    - uses: pre-commit/action@v3.0.1
@@ -3,44 +3,6 @@
 
 Read ahead for some pitfalls, counter-intuitive behavior, and sharp edges that we had to introduce in order to make this work.
 
-.. _tokens:
-
-Token management
-----------------
-
-The compiler behind JAX, XLA, is not aware of the fact that MPI function calls such as :func:`~mpi4jax.send` or :func:`~mpi4jax.recv` must be performed in a specific order (in jargon, that they have *side-effects*). It is therefore free to reorder those calls. Reordering of MPI calls usually leads to deadlocks, e.g. when both processes end up receiving before sending (instead of send-receive, receive-send).
-
-*Tokens* are JAX's way to ensure that XLA does not re-order statements with side effects by injecting a fake data dependency between them.
-
-This means that you *have* to use proper token management to prevent reordering from occurring. Every communication primitive in ``mpi4jax`` returns a token as the last return object, which you have to pass to subsequent primitives within the same JIT block, like this:
-
-.. code:: python
-
-    # DO NOT DO THIS
-    mpi4jax.send(arr, comm=comm)
-    new_arr, _ = mpi4jax.recv(arr, comm=comm)
-
-    # INSTEAD, DO THIS
-    token = mpi4jax.send(arr, comm=comm)
-    new_arr, token = mpi4jax.recv(arr, comm=comm, token=token)
-
-Those functions will then be executed in the same order as the sequence of tokens, from first to last.
-
-Automatic token management
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-An alternative to manual token management is to use the primitives from :mod:`mpi4jax.notoken`, which automatically manage tokens for you. For example, the following code is equivalent to the previous example:
-
-.. code:: python
-
-    import mpi4jax.notoken
-
-    mpi4jax.notoken.send(arr, comm=comm)
-    new_arr = mpi4jax.notoken.recv(arr, comm=comm)
-
-
-This will likely become the default behavior in the future.
-
 No in-place operations in JAX
 -----------------------------
 
 
@@ -42,7 +42,7 @@ The result is an array full of the value 6, because each process adds its rank t
 Basic example: sending and receiving
 ------------------------------------
 
-``mpi4jax`` can of course also send and receive data without performing an operation on it. For this, you can use :func:`~mpi4jax.send` and :func:`~mpi4jax.recv`:
+``mpi4jax`` can send and receive data without performing an operation on it. For this, you can use :func:`~mpi4jax.send` and :func:`~mpi4jax.recv`:
 
 .. _example_2:
 
@@ -64,12 +64,12 @@ Basic example: sending and receiving
         # note: this could also use mpi4jax.sendrecv
         if rank == 0:
             # send, then receive
-            token = mpi4jax.send(arr, dest=1, comm=comm)
-            other_arr, token = mpi4jax.recv(arr, source=1, comm=comm, token=token)
+            mpi4jax.send(arr, dest=1, comm=comm)
+            other_arr = mpi4jax.recv(arr, source=1, comm=comm)
         else:
             # receive, then send
-            other_arr, token = mpi4jax.recv(arr, source=0, comm=comm)
-            token = mpi4jax.send(arr, dest=0, comm=comm, token=token)
+            other_arr = mpi4jax.recv(arr, source=0, comm=comm)
+            mpi4jax.send(arr, dest=0, comm=comm)
 
         return other_arr
 
@@ -91,18 +91,3 @@ Executing this shows that each process has received the data from the other proc
      [1. 1. 1.]]
 
 For operations like this, the correct order of the :func:`~mpi4jax.send` / :func:`~mpi4jax.recv` calls is critical to prevent the program from deadlocking (e.g. when both processes wait for data at the same time).
-
-In ``mpi4jax``, we enforce order of execution through *tokens*. In :ref:`the example code <example_2>`, you can see this behavior e.g. in the following lines:
-
-.. code:: python
-
-    token = mpi4jax.send(arr, dest=1, comm=comm)
-    other_arr, token = mpi4jax.recv(arr, source=1, comm=comm, token=token)
-
-The first call to :func:`~mpi4jax.send` returns a token, which we then pass to :func:`~mpi4jax.recv`. :func:`~mpi4jax.recv` *also* returns a new token that we could pass to subsequent communication primitives.
-
-Because of the nature of JAX, **using tokens to enforce order is not optional.** If you do not use correct token management, you will experience deadlocks and crashes.
-
-.. seealso::
-
-    For more information on tokens, see :ref:`tokens`.
@@ -162,16 +162,15 @@ def get_initial_conditions():
     u0_local = u0_global[local_slice]
     v0_local = v0_global[local_slice]
 
-    token = jax.lax.create_token()
-    h0_local, token = enforce_boundaries(h0_local, "h", token)
-    u0_local, token = enforce_boundaries(u0_local, "u", token)
-    v0_local, token = enforce_boundaries(v0_local, "v", token)
+    h0_local = enforce_boundaries(h0_local, "h")
+    u0_local = enforce_boundaries(u0_local, "u")
+    v0_local = enforce_boundaries(v0_local, "v")
 
     return h0_local, u0_local, v0_local
 
 
 @partial(jax.jit, static_argnums=(1,))
-def enforce_boundaries(arr, grid, token=None):
+def enforce_boundaries(arr, grid):
     """Handle boundary exchange between processors.
 
     This is where mpi4jax comes in!
@@ -222,9 +221,6 @@ def enforce_boundaries(arr, grid, token=None):
         if proc_idx[1] == nproc_x - 1:
             proc_neighbors["east"] = (proc_idx[0], 0)
 
-    if token is None:
-        token = jax.lax.create_token()
-
     for send_dir, recv_dir in zip(send_order, recv_order):
         send_proc = proc_neighbors[send_dir]
         recv_proc = proc_neighbors[recv_dir]
@@ -245,20 +241,17 @@ def enforce_boundaries(arr, grid, token=None):
         send_arr = arr[send_idx]
 
         if send_proc is None:
-            recv_arr, token = mpi4jax.recv(
-                recv_arr, source=recv_proc, comm=mpi_comm, token=token
-            )
+            recv_arr = mpi4jax.recv(recv_arr, source=recv_proc, comm=mpi_comm)
             arr = arr.at[recv_idx].set(recv_arr)
         elif recv_proc is None:
-            token = mpi4jax.send(send_arr, dest=send_proc, comm=mpi_comm, token=token)
+            mpi4jax.send(send_arr, dest=send_proc, comm=mpi_comm)
         else:
-            recv_arr, token = mpi4jax.sendrecv(
+            recv_arr = mpi4jax.sendrecv(
                 send_arr,
                 recv_arr,
                 source=recv_proc,
                 dest=send_proc,
                 comm=mpi_comm,
-                token=token,
             )
             arr = arr.at[recv_idx].set(recv_arr)
 
@@ -268,7 +261,7 @@ def enforce_boundaries(arr, grid, token=None):
     if grid == "v" and proc_idx[0] == nproc_y - 1:
         arr = arr.at[-2, :].set(0.0)
 
-    return arr, token
+    return arr
 
 
 ModelState = namedtuple("ModelState", "h, u, v, dh, du, dv")
@@ -280,20 +273,18 @@ def shallow_water_step(state, is_first_step):
 
     Returns modified model state.
     """
-    token = jax.lax.create_token()
-
     h, u, v, dh, du, dv = state
 
     hc = jnp.pad(h[1:-1, 1:-1], 1, "edge")
-    hc, token = enforce_boundaries(hc, "h", token)
+    hc = enforce_boundaries(hc, "h")
 
     fe = jnp.empty_like(u)
     fn = jnp.empty_like(u)
 
     fe = fe.at[1:-1, 1:-1].set(0.5 * (hc[1:-1, 1:-1] + hc[1:-1, 2:]) * u[1:-1, 1:-1])
     fn = fn.at[1:-1, 1:-1].set(0.5 * (hc[1:-1, 1:-1] + hc[2:, 1:-1]) * v[1:-1, 1:-1])
-    fe, token = enforce_boundaries(fe, "u", token)
-    fn, token = enforce_boundaries(fn, "v", token)
+    fe = enforce_boundaries(fe, "u")
+    fn = enforce_boundaries(fn, "v")
 
     dh_new = dh.at[1:-1, 1:-1].set(
         -(fe[1:-1, 1:-1] - fe[1:-1, :-2]) / dx - (fn[1:-1, 1:-1] - fn[:-2, 1:-1]) / dy
@@ -312,7 +303,7 @@ def shallow_water_step(state, is_first_step):
     q = q.at[1:-1, 1:-1].mul(
         1.0 / (0.25 * (hc[1:-1, 1:-1] + hc[1:-1, 2:] + hc[2:, 1:-1] + hc[2:, 2:]))
     )
-    q, token = enforce_boundaries(q, "h", token)
+    q = enforce_boundaries(q, "h")
 
     du_new = du.at[1:-1, 1:-1].set(
         -GRAVITY * (h[1:-1, 2:] - h[1:-1, 1:-1]) / dx
@@ -337,7 +328,7 @@ def shallow_water_step(state, is_first_step):
             + 0.5 * (v[1:-1, 1:-1] ** 2 + v[:-2, 1:-1] ** 2)
         )
     )
-    ke, token = enforce_boundaries(ke, "h", token)
+    ke = enforce_boundaries(ke, "h")
 
     du_new = du_new.at[1:-1, 1:-1].add(-(ke[1:-1, 2:] - ke[1:-1, 1:-1]) / dx)
     dv_new = dv_new.at[1:-1, 1:-1].add(-(ke[2:, 1:-1] - ke[1:-1, 1:-1]) / dy)
@@ -369,9 +360,9 @@ def shallow_water_step(state, is_first_step):
             )
         )
 
-    h, token = enforce_boundaries(h, "h", token)
-    u, token = enforce_boundaries(u, "u", token)
-    v, token = enforce_boundaries(v, "v", token)
+    h = enforce_boundaries(h, "h")
+    u = enforce_boundaries(u, "u")
+    v = enforce_boundaries(v, "v")
 
     if LATERAL_VISCOSITY > 0:
         # lateral friction
@@ -381,8 +372,8 @@ def shallow_water_step(state, is_first_step):
         fn = fn.at[1:-1, 1:-1].set(
             LATERAL_VISCOSITY * (u[2:, 1:-1] - u[1:-1, 1:-1]) / dy
         )
-        fe, token = enforce_boundaries(fe, "u", token)
-        fn, token = enforce_boundaries(fn, "v", token)
+        fe = enforce_boundaries(fe, "u")
+        fn = enforce_boundaries(fn, "v")
 
         u = u.at[1:-1, 1:-1].add(
             dt
@@ -398,8 +389,8 @@ def shallow_water_step(state, is_first_step):
         fn = fn.at[1:-1, 1:-1].set(
             LATERAL_VISCOSITY * (v[2:, 1:-1] - u[1:-1, 1:-1]) / dy
         )
-        fe, token = enforce_boundaries(fe, "u", token)
-        fn, token = enforce_boundaries(fn, "v", token)
+        fe = enforce_boundaries(fe, "u")
+        fn = enforce_boundaries(fn, "v")
 
         v = v.at[1:-1, 1:-1].add(
             dt