Sharding integration is much slower than pmap

This is a follow up to https://github.com/patrick-kidger/diffrax/issues/403, where when I do the equinox sharding design, I see it is substantially slower than a pmap (or even just a vmap on 1 device).

Here is a MVC that I run on my laptop with 10 devices (10 cpus).

import os
import multiprocessing

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count={}".format(
    multiprocessing.cpu_count()
)

import jax
import jax.numpy as jnp
import jax.experimental.mesh_utils as mesh_utils
import equinox as eqx
import matplotlib.pyplot as plt
from diffrax import *
import optax

def f(t, y, args):
    return jnp.sin(t) + args["theta"] * y

def g(t, y, args):
    return 0.1 * jnp.eye(1)

t0 = 0.
t1 = 1.
dt0 = 0.05
diffusion_shape = jax.ShapeDtypeStruct((1,), "float32")
solver, cont = Heun(), PIDController(1e-3, 1e-6)
ts = jnp.linspace(t0, t1, 100)

def solve(init, key, args):
    control = VirtualBrownianTree(
        t0=t0,
        t1=t1,
        tol=dt0 / 2,
        shape=diffusion_shape,
        key=key,
    )
    vf = ODETerm(f)
    cvf = ControlTerm(g, control)
    terms = MultiTerm(vf, cvf)
    saving = SaveAt(
        ts=ts
    )
    sol = diffeqsolve(
        terms,
        solver,
        y0=init,
        t0=t0,
        t1=t1,
        dt0=dt0,
        args=args,
        saveat=saving,
        stepsize_controller=cont,
        max_steps=10_000,
    )
    return sol.ys

batch_size = 300
inits = 0.1 * jnp.ones((batch_size, 1))
keys = jax.random.split(jax.random.PRNGKey(0), batch_size)
args = {"theta": 0.1}

num_devices = len(jax.devices())
devices = mesh_utils.create_device_mesh((num_devices, 1))
sharding = jax.sharding.PositionalSharding(devices)
replicated = sharding.replicate()

inits_pmap = inits.reshape(num_devices, batch_size // num_devices, *inits.shape[1:])
keys_pmap = keys.reshape(num_devices, batch_size // num_devices, *keys.shape[1:])

args_shard = eqx.filter_shard(args, replicated)
x, y = eqx.filter_shard((inits, keys), sharding)
fn = eqx.filter_jit(eqx.filter_vmap(solve, in_axes=(0, 0, None)), donate="all")
_ = fn(x, y, args_shard).block_until_ready()
_ = eqx.filter_pmap(fn, in_axes=(0, 0, None))(inits_pmap, keys_pmap, args).block_until_ready()

When I do

%%timeit
_ = fn(x, y, args_shard).block_until_ready()

I see 377 ms ± 19.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

and when I do

%%timeit
_ = eqx.filter_pmap(fn, in_axes=(0, 0, None))(inits_pmap, keys_pmap, args).block_until_ready()

I see 10 ms ± 171 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

I thought it could be something with scoping, but if I move everything inside the solve function it has no impact. This just slightly more complex than the pseudo code I described in the previous issue, but seems to be quite slow. Is there an issue with my sharding approach or is there something more complex going on here?

Run the provided multi-device MVC and compare the two timed entry points: fn(x, y, args_shard) using filter_shard versus filter_pmap(fn, in_axes=(0, 0, None)). Check whether the difference remains when changing the device count and batching setup. Done means identifying the source of the performance gap or documenting a confirmed limitation in this sharding approach.