TypeError: No constant handler for type: <class 'jax.core.Token'>
Details can be found here.
https://github.com/mpi4jax/mpi4jax/pull/111

diff --git a/jax/interpreters/mlir.py b/jax/interpreters/mlir.py
index b7aafaf07..771ad7d4e 100644
--- a/jax/interpreters/mlir.py
+++ b/jax/interpreters/mlir.py
@@ -258,6 +258,8 @@ def _device_array_constant_handler(val, canonicalize_types):
 for t in device_array.device_array_types:
   register_constant_handler(t, _device_array_constant_handler)
 
+register_constant_handler(
+  core.Token, lambda _, __: [mhlo.CreateTokenOp(mhlo.TokenType.get())])
 
 # Source locations
 
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 7ab78d7df..606418958 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -374,6 +374,8 @@ for t in device_array.device_array_types:
   register_constant_handler(t, _device_array_constant_handler)
 
 
+register_constant_handler(core.Token, lambda c, _, __: [xops.CreateToken(c)])
+
 # TODO(mattjj): try to remove this canonicalize_dtype stuff
 def canonicalize_dtype(x):
   typ = type(x)
diff --git a/tests/api_test.py b/tests/api_test.py
index 78a33c3e9..c4d012b5a 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2959,6 +2959,11 @@ class APITest(jtu.JaxTestCase):
     x = jax.jit(jax.lax.create_token)(1.0)
     self.assertIsInstance(x, jax.core.Token)
 
+  def test_jit_capturing_token(self):
+    tok = jax.core.token
+    _, y = jax.jit(lambda x: (x + 2, tok))(7)
+    self.assertIsInstance(y, jax.core.Token)
+
   def test_leak_checker_catches_a_jit_leak(self):
     with jax.checking_leaks():
       lst = []
['jax/interpreters/mlir.py', 'jax/interpreters/xla.py', 'tests/api_test.py']