Support for Sharding 2.0 2.0

As a follow up to #735 I wanted to report/pose another issue with Jax sharding and diffrax. While passing the mesh to the shard_map and setting check_vma=False works for the example, automatically inferring the mesh and unsetting the VMA check fails with the following error:

import functools as ft

import jax
import jax.numpy as jnp
from diffrax import ODETerm, Tsit5, diffeqsolve
from jax import shard_map, sharding
from jax.experimental import mesh_utils

mesh = sharding.Mesh(mesh_utils.create_device_mesh((2,)), ["i"])
spec = sharding.PartitionSpec("i")

with jax.set_mesh(mesh):

    @jax.jit
    @ft.partial(shard_map, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False)
    @jax.vmap
    def run(y0):
        term = ODETerm(lambda t, y, args: -y)
        solver = Tsit5()
        t0 = 0
        t1 = 1
        dt0 = 0.1
        sol = diffeqsolve(term, solver, t0, t1, dt0, y0)
        return sol.ys

    y0 = jnp.array([10.0, 10.0])
    y0 = jax.device_put(y0, spec)
    out = run(y0)
    print(out)
    jax.debug.visualize_array_sharding(out)

The stack trace looks as follows:

Traceback (most recent call last):
  File "/workspaces/gfdm-jax/scripts/test.py", line 28, in <module>
    out = run(y0)
          ^^^^^^^
  File "/workspaces/gfdm-jax/scripts/test.py", line 23, in run
    sol = diffeqsolve(term, solver, t0, t1, dt0, y0)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_integrate.py", line 1456, in diffeqsolve
    final_state, aux_stats = adjoint.loop(
                             ^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/equinox/_module/_prebuilt.py", line 34, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_adjoint.py", line 299, in loop
    final_state = self._loop(
                  ^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_integrate.py", line 666, in loop
    _, traced_jump, traced_result = eqx.filter_eval_shape(body_fun_aux, init_state)
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_integrate.py", line 373, in body_fun_aux
    (y, y_error, dense_info, solver_state, solver_result) = solver.step(
                                                            ^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/equinox/_module/_prebuilt.py", line 34, in __call__
    return self.__func__(self.__self__, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_solver/runge_kutta.py", line 1194, in step
    y_error = jtu.tree_map(
              ^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_solver/runge_kutta.py", line 1195, in <lambda>
    lambda _y_error: jnp.where(is_okay(result), _y_error, jnp.inf),
                               ^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/diffrax/_solution.py", line 54, in is_okay
    return is_successful(result) | is_event(result)
           ~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py", line 609, in deferring_binary_op
    return binary_op(*args)
           ^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/ufunc_api.py", line 182, in __call__
    return call(*args)
           ^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/sharding.py", line 197, in shard_shape
    return _common_shard_shape(self, global_shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/sharding.py", line 64, in _common_shard_shape
    partitions, _ = get_num_ways_dim_sharded(hlo_sharding)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/workspaces/gfdm-jax/.venv/lib/python3.12/site-packages/jax/_src/op_shardings.py", line 30, in get_num_ways_dim_sharded
    assert not hlo_sharding.is_manual()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

I tried working around that, but ended up hacking my code to not execute the function under the mesh context. Since the context manager is the recommended way to handle meshes currently, I think that more people will face this bug though.

Start by running the provided reproduction under the jax.set_mesh context, then trace the failure through diffrax/_solver/runge_kutta.py and diffrax/_solution.py as shown in the stack trace. Done means the example with automatically inferred mesh and check_vma enabled completes without the JAX manual-sharding AssertionError.