Repository: google/jax (python)

Issue:
[Feature Request] mish activation function
Hi, @mattjj @froystig I would like to contribute by adding  ```mish``` activation in  Jax function. It's a  Self Regularized Non-Monotonic Neural Activation Function.

```mish``` =  ```x * tanh(softplus(x))```
paper: https://arxiv.org/abs/1908.08681

Pull request description:
Add jax.nn.mish.
Fixes #4523

Patch:
diff --git a/.github/workflows/wheel_win_x64.yml b/.github/workflows/wheel_win_x64.yml
index 33457ab3d..a6d310fb1 100644
--- a/.github/workflows/wheel_win_x64.yml
+++ b/.github/workflows/wheel_win_x64.yml
@@ -37,7 +37,7 @@ jobs:
         run: |
           python -m pip install -r build/test-requirements.txt
           python -m pip uninstall -y matplotlib
-          python -m pip install --pre --upgrade numpy==2.0.0rc1 scipy==1.13.0
+          python -m pip install --pre --upgrade numpy==2.0.0rc1 scipy==1.13.0rc1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
           python.exe build\build.py --bazel_options=--color=yes --verbose
 
diff --git a/.github/workflows/windows_ci.yml b/.github/workflows/windows_ci.yml
index 811f09512..234187e38 100644
--- a/.github/workflows/windows_ci.yml
+++ b/.github/workflows/windows_ci.yml
@@ -47,7 +47,7 @@ jobs:
           cd jax
           python -m pip install -r build/test-requirements.txt
           python -m pip uninstall -y matplotlib
-          python -m pip install --pre --upgrade numpy==2.0.0rc1 scipy==1.13.0
+          python -m pip install --pre --upgrade numpy==2.0.0rc1 scipy==1.13.0rc1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
           python.exe build\build.py ('--bazel_options=--override_repository=xla=${{ github.workspace }}\xla' -replace '\\','\\') --bazel_options=--color=yes
 
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1ebbf585d..a683ee83b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -50,7 +50,6 @@ Remember to align the itemized text with the first line of an item within a list
 * Changes
   * JAX now supports CUDA 12.1 or newer only. Support for CUDA 11.8 has been
     dropped.
-  * JAX now supports NumPy 2.0.
 
 ## jax 0.4.25 (Feb 26, 2024)
 
diff --git a/docs/jax.nn.rst b/docs/jax.nn.rst
index 4cb669540..67d80e160 100644
--- a/docs/jax.nn.rst
+++ b/docs/jax.nn.rst
@@ -38,6 +38,7 @@ Activation functions
     gelu
     glu
     squareplus
+    mish
 
 Other functions
 ---------------
diff --git a/docs/multi_process.md b/docs/multi_process.md
index ab4beaf8c..70bead4c4 100644
--- a/docs/multi_process.md
+++ b/docs/multi_process.md
@@ -62,9 +62,6 @@ The API {func}`jax.distributed.initialize` takes several arguments, namely:
     with a port available on that process. Process 0 will start a JAX service
     exposed via that IP address and port, to which the other processes in the
     cluster will connect.
-  * `coordinator_bind_address`: the IP address and port to which the JAX service
-    on process 0 in your cluster will bind. By default, it will bind to all
-    available interfaces using the same port as `coordinator_address`.
   * `num_processes`: the number of processes in the cluster
   * `process_id`: the ID number of this process, in the range `[0 ..
     num_processes)`.
diff --git a/docs/pallas/design.md b/docs/pallas/design.md
index aa67d60c4..815a6c315 100644
--- a/docs/pallas/design.md
+++ b/docs/pallas/design.md
@@ -88,7 +88,6 @@ Because JAX was designed with HLO in mind, the set of JAX primitives closely mir
 Because Pallas was initially designed with Triton in mind, we offer a set of new primitives targeting the Triton programming model. As we’ll show later, we can lower these primitives to Mosaic as well.
 
 #### `pallas.load` and `pallas.store`
-
 `pallas.load` and `pallas.store` are primitives that allow loading from memory and storing into memory. Unlike `__getitem__` and `__setitem__` they are more flexible at the cost of being more verbose. Specifically, you can use the `pallas.dynamic_slice` (`pallas.ds` for short) construct (which should maybe be upstreamed into JAX to be used with Ref `__getitem__` and `__setitem__`).
 
 ```python
@@ -115,7 +114,6 @@ def f(x_ref, o_ref):
 Masking is important when doing out-of-bounds loads/stores. The operational semantics of masking can be compiler-determined (if we understand the documentation properly, Triton avoids the read from/write to memory if it’s masked).
 
 #### `pallas.program_id` and `pallas.num_programs`
-
 As we’ll soon see, we’ll be executing the same Pallas kernels many times (either in parallel or in a pipeline depending on the backend). These new primitives tell us “where” we are in the execution of the kernel.
 
 `pallas.program_id` takes in an axis argument, which tells us which index in an axis of a multidimensional grid this kernel is currently executing in (analogous to `threadId` from CUDA programming or `lax.axis_index` in `jax.pmap`). Note that we are currently borrowing the “program” terminology from Triton and in the future we might want to change it to something more familiar to JAX users.
@@ -253,7 +251,7 @@ In this example, we compute tiles of the output by doing an unrolled accumulatio
 ```python
 def matmul_kernel(x_ref, y_ref, o_ref, *, activation, block_k):
   acc = jnp.zeros((x_ref.shape[0], x_ref.shape[1]), jnp.float32)
-  for k in range(x_ref.shape[1] // block_k):
+  for k in range(x_ref.shape[1] // block_k)
     x = x_ref[:, k*block_k:(k+1)*block_k]
     y = y_ref[k*block_k:(k+1)*block_k, :]
     acc += x @ y
@@ -269,12 +267,10 @@ def matmul(x, y, *, block_shape, activation):
       partial(matmul_kernel, block_k=block_k, activation=activation),
       out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1],), jnp.float32),
       in_specs=[
-        pl.BlockSpec(lambda i, j: (i, 0), (block_m, x.shape[1])),
-        pl.BlockSpec(lambda i, j: (0, j), (y.shape[0], block_n))
+        pl.BlockSpec(lambda i, j:, (i, 0), (block_m, x.shape[1])),
+        pl.BlockSpec(lambda i, j:, (0, j), (y.shape[0], block_n))
       ],
-      out_specs=pl.BlockSpec(lambda i, j: (i, j), (block_m, block_n)),
-      grid=(4, 4),
-  )
+      out_specs=pl.BlockSpec(lambda i, j: (i, j), (block_m, block_n))
   return fused_matmul(x, y)
 
 z = matmul(x, y, block_shape=block_shape, activation=jax.nn.gelu)
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index 0942d04ab..b0b66b2f4 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -244,7 +244,7 @@ def _check_shape_dtype(shape_dtype):
   dt = np.dtype(shape_dtype.dtype)
   if dtypes.canonicalize_dtype(dt) != dt:
     raise ValueError(
-        "result_shape_dtypes cannot specify 64-bit types when `jax_enable_x64` is disabled")
+        "Cannot return 64-bit values when `jax_enable_x64` is disabled")
 
 
 def pure_callback(
diff --git a/jax/_src/distributed.py b/jax/_src/distributed.py
index 508f124a3..63464bb0e 100644
--- a/jax/_src/distributed.py
+++ b/jax/_src/distributed.py
@@ -41,8 +41,7 @@ class State:
                  num_processes: int | None = None,
                  process_id: int | None = None,
                  local_device_ids: int | Sequence[int] | None = None,
-                 initialization_timeout: int = 300,
-                 coordinator_bind_address: str | None = None):
+                 initialization_timeout: int = 300):
     coordinator_address = (coordinator_address or
                            os.environ.get('JAX_COORDINATOR_ADDRESS', None))
     if isinstance(local_device_ids, int):
@@ -67,15 +66,6 @@ class State:
 
     self.coordinator_address = coordinator_address
 
-    # The default value of [::]:port tells the coordinator to bind to all
-    # available addresses on the same port as coordinator_address.
-    default_coordinator_bind_address = '[::]:' + coordinator_address.rsplit(':', 1)[1]
-    coordinator_bind_address = (coordinator_bind_address or
-                                os.environ.get('JAX_COORDINATOR_BIND_ADDRESS',
-                                               default_coordinator_bind_address))
-    if coordinator_bind_address is None:
-      raise ValueError('coordinator_bind_address should be defined.')
-
     if local_device_ids:
       visible_devices = ','.join(str(x) for x in local_device_ids) # type: ignore[union-attr]
       logger.info('JAX distributed initialized with visible devices: %s', visible_devices)
@@ -89,7 +79,7 @@ class State:
         raise RuntimeError('distributed.initialize should only be called once.')
       logger.info('Starting JAX distributed service on %s', coordinator_address)
       self.service = xla_extension.get_distributed_runtime_service(
-          coordinator_bind_address, num_processes)
+          coordinator_address, num_processes)
 
     self.num_processes = num_processes
 
@@ -128,8 +118,7 @@ def initialize(coordinator_address: str | None = None,
                num_processes: int | None = None,
                process_id: int | None = None,
                local_device_ids: int | Sequence[int] | None = None,
-               initialization_timeout: int = 300,
-               coordinator_bind_address: str | None = None):
+               initialization_timeout: int = 300):
   """Initializes the JAX distributed system.
 
   Calling :func:`~jax.distributed.initialize` prepares JAX for execution on
@@ -167,11 +156,6 @@ def initialize(coordinator_address: str | None = None,
     initialization_timeout: Time period (in seconds) for which connection will
       be retried. If the initialization takes more than the timeout specified,
       the initialization will error. Defaults to 300 secs i.e. 5 mins.
-    coordinator_bind_address: the address and port to which the coordinator service
-      on process `0` should bind. If this is not specified, the default is to bind to
-      all available addresses on the same port as ``coordinator_address``. On systems
-      that have multiple network interfaces per node it may be insufficient to only
-      have the coordinator service listen on one address/interface.
 
   Raises:
     RuntimeError: If :func:`~jax.distributed.initialize` is called more than once.
@@ -194,7 +178,7 @@ def initialize(coordinator_address: str | None = None,
     raise RuntimeError("jax.distributed.initialize() must be called before "
                         "any JAX computations are executed.")
   global_state.initialize(coordinator_address, num_processes, process_id,
-                          local_device_ids, initialization_timeout, coordinator_bind_address)
+                          local_device_ids, initialization_timeout)
   atexit.register(shutdown)
 
 
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 239559a47..d2572b5d2 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -2469,12 +2469,16 @@ def emit_python_callback(
           "Mismatched number of outputs from callback. "
           "Expected: {}, Actual: {}".format(len(result_avals), len(out_vals)))
     # Handle Python literals, and custom arrays, e.g., tf.Tensor.
-    out_vals = tuple(xla.canonicalize_dtype(np.asarray(a)) for a in out_vals)
+    out_vals = tuple(np.asarray(a) for a in out_vals)
     for i, (out_val, out_aval) in enumerate(zip(out_vals, result_avals)):
       if out_val.shape != out_aval.shape:
         raise RuntimeError(
             f"Incorrect output shape for return value #{i}: "
             f"Expected: {out_aval.shape}, Actual: {out_val.shape}")
+      if out_val.dtype != dtypes.canonicalize_dtype(out_val.dtype):
+        raise RuntimeError(
+            "Cannot return 64-bit values when `jax_enable_x64` is disabled. "
+            f"Actual: {out_val.dtype}")
       if out_val.dtype != out_aval.dtype:
         raise RuntimeError(
             f"Incorrect output dtype for return value #{i}: "
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
index 2a436e315..c9e0b3db7 100644
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -199,6 +199,29 @@ def silu(x: ArrayLike) -> Array:
 
 swish = silu
 
+@jax.jit
+def mish(x: ArrayLike) -> Array:
+  r"""Mish activation function.
+
+  Computes the element-wise function:
+
+  .. math::
+    \mathrm{mish}(x) = x \cdot \mathrm{tanh}(\mathrm{softplus}(x))
+
+  For more information, see
+  `Mish: A Self Regularized Non-Monotonic Activation Function
+  <https://arxiv.org/abs/1908.08681>`_.
+
+  Args:
+    x : input array
+
+  Returns:
+    An array.
+  """
+  numpy_util.check_arraylike("mish", x)
+  x_arr = jnp.asarray(x)
+  return x_arr * jnp.tanh(softplus(x_arr))
+
 @jax.jit
 def log_sigmoid(x: ArrayLike) -> Array:
   r"""Log-sigmoid activation function.
@@ -314,7 +337,7 @@ def celu(x: ArrayLike, alpha: ArrayLike = 1.0) -> Array:
 
   For more information, see
   `Continuously Differentiable Exponential Linear Units
-  <https://arxiv.org/pdf/1704.07483.pdf>`_.
+  <https://arxiv.org/abs/1704.07483>`_.
 
   Args:
     x : input array
@@ -342,7 +365,7 @@ def selu(x: ArrayLike) -> Array:
 
   For more information, see
   `Self-Normalizing Neural Networks
-  <https://papers.nips.cc/paper/6698-self-normalizing-neural-networks.pdf>`_.
+  <https://arxiv.org/abs/1706.02515>`_.
 
   Args:
     x : input array
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 826251c82..0c2640186 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -3443,11 +3443,8 @@ def vdot(
              preferred_element_type=preferred_element_type)
 
 
-@util.implements(
-    getattr(np, "vecdot", None), lax_description=_PRECISION_DOC,
-    extra_params=_DOT_PREFERRED_ELEMENT_TYPE_DESCRIPTION,
-    # TODO(phawkins): numpy.vecdot doesn't have a __module__ attribute.
-    module="numpy")
+@util.implements(getattr(np, "vecdot", None), lax_description=_PRECISION_DOC,
+                 extra_params=_DOT_PREFERRED_ELEMENT_TYPE_DESCRIPTION)
 def vecdot(x1: ArrayLike, x2: ArrayLike, /, *, axis: int = -1,
            precision: PrecisionLike = None,
            preferred_element_type: DTypeLike | None = None) -> Array:
diff --git a/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py b/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py
index a8f72e6eb..badb292aa 100644
--- a/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py
+++ b/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py
@@ -93,18 +93,14 @@ def paged_flash_attention_kernel(
     pages_per_sequence: int,
     mask_value: float,
     megacore_mode: str,
-    program_ids=(),
 ):
   """Pallas kernel for paged attention."""
-  if program_ids:
-    core_index, b, h, i = program_ids
-  else:
-    core_index, b, h, i = (
-        pl.program_id(0),
-        pl.program_id(1),
-        pl.program_id(2),
-        pl.program_id(3),
-    )
+  core_index, b, h, i = (
+      pl.program_id(0),
+      pl.program_id(1),
+      pl.program_id(2),
+      pl.program_id(3),
+  )
   num_kv_heads, _, page_size, _ = k_pages_hbm_ref.shape
   bk = page_size * pages_per_compute_block
   num_cores = pl.num_programs(0)
@@ -118,35 +114,25 @@ def paged_flash_attention_kernel(
   b = b * b_step + b_start
   length = lengths_ref[b]
 
-  def compute_block_indices(b, h, i):
-
-    def advance_b():
-      next_b = b + b_step
-
-      def advance_to_next_non_zero_length():
-        next_next_b = next_b + b_step
-        return lax.fori_loop(
-            lax.div(next_next_b, b_step),
-            lax.div(batch_size, b_step),
-            lambda _, b: jnp.where(lengths_ref[b] == 0, b + b_step, b),
-            next_next_b,
-        )
-
-      return (
-          lax.cond(
-              jnp.logical_and(next_b < batch_size, lengths_ref[next_b] == 0),
-              advance_to_next_non_zero_length,
-              lambda: next_b,
-          ),
-          h_start,
-          0,
-      )
-
-    def advance_h():
-      next_h = h + h_step
-      return lax.cond(next_h < num_kv_heads, lambda: (b, next_h, 0), advance_b)
+  def advance_to_next_non_zero_length(b):
+    return lax.fori_loop(
+        lax.div(b, b_step),
+        lax.div(batch_size, b_step),
+        lambda _, b: jnp.where(lengths_ref[b] == 0, b + b_step, b),
+        b,
+    )
 
-    return lax.cond(i * bk < lengths_ref[b], lambda: (b, h, i), advance_h)
+  def compute_block_indices(b, h, i):
+    length = lengths_ref[b]
+    not_done = i * bk < length
+    i_next = jnp.where(not_done, i, 0)
+    h_next = jnp.where(not_done, h, h + h_step)
+    is_last_head = h_next >= num_kv_heads
+    h_next = jnp.where(is_last_head, h_start, h_next)
+    b_next = jnp.where(
+        is_last_head, advance_to_next_non_zero_length(b + b_step), b
+    )
+    return b_next, h_next, i_next
 
   def create_kv_async_copy_descriptors(b, h, i, buffer_index):
     page_offset = b * pages_per_sequence + i * pages_per_compute_block
@@ -233,64 +219,12 @@ def paged_flash_attention_kernel(
     step_ref[0] = step + 1
 
 
-def paged_flash_attention_kernel_inline_seq_dim(
-    lengths_ref,
-    page_indices_ref,
-    buffer_index_ref,
-    step_ref,
-    q_ref,
-    k_pages_hbm_ref,
-    v_pages_hbm_ref,
-    o_ref,
-    m_ref,
-    l_ref,
-    k_vmem_buffer,
-    v_vmem_buffer,
-    sem,
-    *,
-    batch_size: int,
-    pages_per_compute_block: int,
-    pages_per_sequence: int,
-    mask_value: float,
-    megacore_mode: str,
-):
-  core_index, b, h = pl.program_id(0), pl.program_id(1), pl.program_id(2)
-
-  def body(i, _):
-    paged_flash_attention_kernel(
-        lengths_ref,
-        page_indices_ref,
-        buffer_index_ref,
-        step_ref,
-        q_ref,
-        k_pages_hbm_ref,
-        v_pages_hbm_ref,
-        o_ref,
-        m_ref,
-        l_ref,
-        k_vmem_buffer,
-        v_vmem_buffer,
-        sem,
-        batch_size=batch_size,
-        pages_per_compute_block=pages_per_compute_block,
-        pages_per_sequence=pages_per_sequence,
-        mask_value=mask_value,
-        megacore_mode=megacore_mode,
-        program_ids=(core_index, b, h, i),
-    )
-    return ()
-
-  bk = pages_per_compute_block * k_pages_hbm_ref.shape[-2]
-  lax.fori_loop(0, lax.div(lengths_ref[b] + bk - 1, bk), body, ())
-
-
 @functools.partial(
     jax.jit,
     static_argnames=[
         "pages_per_compute_block",
         "mask_value",
         "megacore_mode",
-        "inline_seq_dim",
     ],
 )
 def paged_attention(
@@ -303,7 +237,6 @@ def paged_attention(
     mask_value: float = DEFAULT_MASK_VALUE,
     pages_per_compute_block: int,
     megacore_mode: Optional[str] = None,
-    inline_seq_dim: bool = True,
 ) -> tuple[jax.Array, tuple[jax.Array, jax.Array]]:
   """Paged grouped query attention.
 
@@ -328,8 +261,6 @@ def paged_attention(
         divisible by 2.
       * batch: megacore parallelism on batch dimension; requires batch divisible
         by 2.
-    inline_seq_dim: whether to fuse kernel instances along the sequence dim into
-      one kernel.
 
   Returns:
     The output of attention([batch_size, num_heads, head_dim]).
@@ -421,31 +352,9 @@ def paged_attention(
       )
     q_dtype_for_kernel_launch = q.dtype
 
-  if inline_seq_dim:
-    kernel = paged_flash_attention_kernel_inline_seq_dim
-    grid = (
-        num_cores,
-        batch_size // num_cores if megacore_mode == "batch" else batch_size,
-        num_kv_heads // num_cores
-        if megacore_mode == "kv_head"
-        else num_kv_heads,
-    )  # type: ignore
-    dimension_sematics = ("parallel", "arbitrary", "arbitrary")  # type: ignore
-  else:
-    kernel = paged_flash_attention_kernel
-    grid = (
-        num_cores,
-        batch_size // num_cores if megacore_mode == "batch" else batch_size,
-        num_kv_heads // num_cores
-        if megacore_mode == "kv_head"
-        else num_kv_heads,
-        pages_per_sequence // pages_per_compute_block,
-    )  # type: ignore
-    dimension_sematics = ("parallel", "arbitrary", "arbitrary", "arbitrary")  # type: ignore
-
   out, _, _ = pl.pallas_call(
       functools.partial(
-          kernel,
+          paged_flash_attention_kernel,
           pages_per_sequence=pages_per_sequence,
           batch_size=batch_size,
           pages_per_compute_block=pages_per_compute_block,
@@ -464,7 +373,16 @@ def paged_attention(
               q_block_spec,
               q_block_spec,
           ],
-          grid=grid,
+          grid=(
+              num_cores,
+              batch_size // num_cores
+              if megacore_mode == "batch"
+              else batch_size,
+              num_kv_heads // num_cores
+              if megacore_mode == "kv_head"
+              else num_kv_heads,
+              pages_per_sequence // pages_per_compute_block,
+          ),
           scratch_shapes=(
               pltpu.VMEM(
                   (
@@ -488,7 +406,12 @@ def paged_attention(
           ),
       ),
       mosaic_params=dict(
-          dimension_semantics=dimension_sematics,
+          dimension_semantics=(
+              "parallel",
+              "arbitrary",
+              "arbitrary",
+              "arbitrary",
+          )
       ),
       out_shape=[
           jax.ShapeDtypeStruct(q.shape, q_dtype_for_kernel_launch),
diff --git a/jax/nn/__init__.py b/jax/nn/__init__.py
index 0f7e42560..2d4fd848b 100644
--- a/jax/nn/__init__.py
+++ b/jax/nn/__init__.py
@@ -45,6 +45,7 @@ from jax._src.nn.functions import (
   silu as silu,
   swish as swish,
   squareplus as squareplus,
+  mish as mish,
 )
 
 # Deprecations
diff --git a/setup.py b/setup.py
index bcf687dfd..3b465c0ef 100644
--- a/setup.py
+++ b/setup.py
@@ -22,12 +22,12 @@ from setuptools import setup, find_packages
 
 project_name = 'jax'
 
-_current_jaxlib_version = '0.4.26'
+_current_jaxlib_version = '0.4.25'
 # The following should be updated with each new jaxlib release.
 _latest_jaxlib_version_on_pypi = '0.4.25'
 _default_cuda12_cudnn_version = '89'
 _available_cuda12_cudnn_versions = [_default_cuda12_cudnn_version]
-_libtpu_version = '0.1.dev20240403'
+_libtpu_version = '0.1.dev20240224'
 
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(


Test patch:
diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
index e497b168a..63d9597f0 100644
--- a/jax/_src/test_util.py
+++ b/jax/_src/test_util.py
@@ -1746,24 +1746,6 @@ class numpy_with_mpmath:
   def expm1(self, x):
     return x.context.expm1(x)
 
-  def log1p(self, x):
-    ctx = x.context
-    if isinstance(x, ctx.mpc):
-      # Workaround mpmath 1.3 bug in log(+-inf+-infj) evaluation (see mpmath/mpmath#774).
-      # TODO(pearu): remove this function when mpmath 1.4 or newer
-      # will be the required test dependency.
-      if ctx.isinf(x.real) and ctx.isinf(x.imag):
-        pi = ctx.pi
-        if x.real > 0 and x.imag > 0:
-          return ctx.make_mpc((x.real._mpf_, (pi / 4)._mpf_))
-        if x.real > 0 and x.imag < 0:
-          return ctx.make_mpc((x.real._mpf_, (-pi / 4)._mpf_))
-        if x.real < 0 and x.imag < 0:
-          return ctx.make_mpc(((-x.real)._mpf_, (-3 * pi / 4)._mpf_))
-        if x.real < 0 and x.imag > 0:
-          return ctx.make_mpc(((-x.real)._mpf_, (3 * pi / 4)._mpf_))
-    return ctx.log1p(x)
-
   def log2(self, x):
     return x.context.ln(x) / x.context.ln2
 
diff --git a/tests/lax_test.py b/tests/lax_test.py
index fead85aa2..eecb93a9b 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -3666,9 +3666,11 @@ class FunctionAccuracyTest(jtu.JaxTestCase):
         regions_with_inaccuracies_keep('q1', 'q2', 'q3', 'q4', 'neg', 'pos', 'negj', 'posj', 'ninf', 'pinf', 'ninfj', 'pinfj')
 
     elif name == 'arctanh':
-      regions_with_inaccuracies_keep('q1', 'q2', 'q3', 'q4', 'neg', 'pos', 'negj', 'posj', 'ninf', 'pinf', 'ninfj', 'pinfj', 'mpos')
-      # TODO(pearu): after landing openxla/xla#10503, switch to
-      # regions_with_inaccuracies_keep('pos', 'ninf', 'pinf', 'ninfj', 'pinfj', 'mpos')
+      if xla_extension_version < 251:
+        regions_with_inaccuracies_keep('q1', 'q2', 'q3', 'q4', 'neg', 'pos', 'negj', 'posj', 'ninf', 'pinf', 'ninfj', 'pinfj', 'mpos')
+      else:
+        regions_with_inaccuracies_keep('pos', 'ninf', 'pinf', 'ninfj', 'pinfj', 'mpos')
+
     elif name in {'positive', 'negative', 'conjugate', 'sin', 'cos', 'sqrt', 'expm1'}:
       regions_with_inaccuracies.clear()
     else:
diff --git a/tests/nn_test.py b/tests/nn_test.py
index 67890db9e..32e1e78b0 100644
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -91,6 +91,26 @@ class NNFunctionsTest(jtu.JaxTestCase):
   def testSquareplusZero(self, dtype):
     self.assertEqual(dtype(1), nn.squareplus(dtype(0), dtype(4)))
 
+  def testMishGrad(self):
+    check_grads(nn.mish, (1e-8,), order=4,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testMishGradZero(self):
+    check_grads(nn.mish, (0.,), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testMishGradNegInf(self):
+    check_grads(nn.mish, (-float('inf'),), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  def testMishGradNan(self):
+    check_grads(nn.mish, (float('nan'),), order=1,
+                rtol=1e-2 if jtu.test_device_matches(["tpu"]) else None)
+
+  @parameterized.parameters([float] + jtu.dtypes.floating)
+  def testMishZero(self, dtype):
+    self.assertEqual(dtype(0), nn.mish(dtype(0)))
+
   def testReluGrad(self):
     rtol = 1e-2 if jtu.test_device_matches(["tpu"]) else None
     check_grads(nn.relu, (1.,), order=3, rtol=rtol)
@@ -117,6 +137,10 @@ class NNFunctionsTest(jtu.JaxTestCase):
     val = nn.squareplus(1e3)
     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
+  def testMishValue(self):
+    val = nn.mish(1e3)
+    self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
+
   @jtu.skip_on_flag("jax_skip_slow_tests", True)
   def testEluGrad(self):
     check_grads(nn.elu, (1e4,), order=4, eps=1.)
@@ -149,7 +173,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
       (jnp.float32, jnp.bfloat16, jnp.float16),
       (partial(nn.gelu, approximate=False),
        partial(nn.gelu, approximate=True),
-       nn.relu, nn.softplus, nn.sparse_plus, nn.sigmoid, nn.squareplus)))
+       nn.relu, nn.softplus, nn.sparse_plus, nn.sigmoid, nn.squareplus, nn.mish)))
   def testDtypeMatchesInput(self, dtype, fn):
     x = jnp.zeros((), dtype=dtype)
     out = fn(x)
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index 2329fe65e..0cd51631a 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -117,10 +117,10 @@ class PythonCallbackTest(jtu.JaxTestCase):
           lambda x: returned_literal, core.ShapedArray((), expect_dtype), x
       )
 
-    if not config.enable_x64.value and expect_dtype in (np.int64, np.float64):
-      ctx = self.assertRaisesRegex(Exception, "result_shape_dtypes cannot specify 64-bit types")
-    elif config.enable_x64.value and expect_dtype in (np.int32, np.float32):
-      ctx = self.assertRaisesRegex(Exception, "Incorrect output dtype for return value")
+    if not config.enable_x64.value:
+      ctx = self.assertRaisesRegex(Exception, "Cannot return 64-bit values")
+    elif expect_dtype in (np.int32, np.float32):
+      ctx = self.assertRaisesRegex(Exception, "Incorrect output dtype")
     else:
       ctx = contextlib.nullcontext()
 
@@ -247,7 +247,7 @@ class PythonCallbackTest(jtu.JaxTestCase):
       jax.effects_barrier()
 
   @with_pure_and_io_callbacks
-  def test_callback_with_wrong_dtype_outputs(self, *, callback=io_callback_ordered):
+  def test_callback_with_wrong_dtype_outputs(self, *, callback):
 
     def _cb():
       return np.array([1], np.float64)
@@ -257,14 +257,9 @@ class PythonCallbackTest(jtu.JaxTestCase):
       # Calling a function expected a f32 return value but getting f64
       return callback(_cb, core.ShapedArray((1,), np.float32))
 
-    if config.enable_x64.value:
-      ctx = self.assertRaisesRegex(Exception, "Incorrect output dtype for return value")
-    else:
-      ctx = contextlib.nullcontext()
-    with ctx:
-      res = f()
+    with self.assertRaises(RuntimeError):
+      f()
       jax.effects_barrier()
-      self.assertAllClose(res, np.array([1], np.float32))
 
   @with_pure_and_io_callbacks
   def test_callback_with_wrongly_specified_64_bit_dtype(self, *, callback):
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
index f665a7b55..d2d5b7c25 100644
--- a/tests/random_lax_test.py
+++ b/tests/random_lax_test.py
@@ -65,14 +65,6 @@ class LaxRandomTest(jtu.JaxTestCase):
     # whether RBG keys may be involved, but that's no longer exact.
     if config.enable_custom_prng.value and samples.dtype == jnp.bfloat16:
       return
-    # kstest fails for infinities starting in scipy 1.12
-    # (https://github.com/scipy/scipy/issues/20386)
-    # TODO(jakevdp): remove this logic if/when fixed upstream.
-    scipy_version = jtu.parse_version(scipy.__version__)
-    if scipy_version >= (1, 12) and np.issubdtype(samples.dtype, np.floating):
-      samples = np.array(samples, copy=True)
-      samples[np.isposinf(samples)] = 0.01 * np.finfo(samples.dtype).max
-      samples[np.isneginf(samples)] = 0.01 * np.finfo(samples.dtype).min
     self.assertGreater(scipy.stats.kstest(samples, cdf).pvalue, fail_prob)
 
   def _CheckChiSquared(self, samples, pmf, *, pval=None):
@@ -750,6 +742,9 @@ class LaxRandomTest(jtu.JaxTestCase):
   )
   @jtu.skip_on_devices("cpu", "tpu")  # TODO(phawkins): slow compilation times
   def testT(self, df, dtype):
+    scipy_version = jtu.parse_version(scipy.__version__)
+    if scipy_version >= (1, 13):
+      self.skipTest("ks test returns NaN on SciPy 1.13")
     key = lambda: self.make_key(1)
     rand = lambda key, df: random.t(key, df, (10000,), dtype)
     crand = jax.jit(rand)