Raise a runtime error when trying to convert the jax.Array wrapped …

…by `jax.core.Token` to a numpy array, as it is an internal implementation detail and the buffer has XLA token shape.

PiperOrigin-RevId: 632682906

tests/api_test.py

@@ -58,6 +58,7 @@
 from jax._src.interpreters import partial_eval as pe
 from jax._src.lib import xla_client
 from jax._src.lib import xla_extension
+from jax._src.lib import xla_extension_version
 import jax._src.util as jax_util
 from jax.ad_checkpoint import checkpoint_name, checkpoint as new_checkpoint
 import jax.custom_batching
@@ -702,6 +703,14 @@ def test_trivial_computations(self):
     self.assertNotEqual(z3.unsafe_buffer_pointer(), x1.unsafe_buffer_pointer())
     self.assertEqual(z2, 1)
 
+  @unittest.skipIf(xla_extension_version < 264, "jaxlib version too old")
+  def test_print_token_buffer_error(self):
+    token = jax.lax.create_token()
+    with self.assertRaisesRegex(
+        RuntimeError, "Cannot convert a token-shape buffer to a numpy array."
+    ):
+      token._buf._value
+
   def test_trivial_computations_with_tokens(self):
     @jit
     def noop(arr, token):