Adjoints question

I was putting together some tests, when I realized that I'm not sure I fully understand the error bounds on the non recursive (which is what I almost always use) adjoints. Specifically, for a simple diagonal noise SDE, I actually encounter errors for the other adjoints, are these expected or am I doing something wrong?

import jax
import jax.numpy as jnp
import jax.random as jr
import equinox as eqx
import diffrax
import lineax as lx

key = jax.random.PRNGKey(42)
key, subkey = jax.random.split(key)
driftkey, diffusionkey, ykey = jr.split(subkey, 3)
drift_mlp = eqx.nn.MLP(
    in_size=3,
    out_size=3,
    width_size=8,
    depth=2,
    activation=jax.nn.swish,
    final_activation=jnp.tanh,
    key=driftkey,
)
diffusion_mlp = eqx.nn.MLP(
    in_size=3,
    out_size=3,
    width_size=8,
    depth=2,
    activation=jax.nn.swish,
    final_activation=jnp.tanh,
    key=diffusionkey,
)

class Field(eqx.Module):
    force: eqx.nn.MLP

    def __call__(self, t, y, args):
        return self.force(y)

class DiffusionField(eqx.Module):
    force: eqx.nn.MLP

    def __call__(self, t, y, args):
        return lx.DiagonalLinearOperator(self.force(y))

y0 = jr.normal(ykey, (3,))

k1, k2, k3 = jax.random.split(key, 3)
vbt = diffrax.VirtualBrownianTree(0.3, 9.5, 1e-4, (3,), k1, levy_area=diffrax.SpaceTimeLevyArea)
vbt_terms = diffrax.MultiTerm(
    diffrax.ODETerm(Field(drift_mlp)),
    diffrax.ControlTerm(DiffusionField(diffusion_mlp), vbt),
)
solver = diffrax.GeneralShARK()
y0_args_term0 = (y0, None, vbt_terms)

def _run(y0__args__term, saveat, adjoint):
    y0_, args, term = y0__args__term
    ys = diffrax.diffeqsolve(
        term,
        solver,
        0.3,
        9.5,
        0.1,
        y0_,
        args,
        saveat=saveat,
        adjoint=adjoint,
    ).ys
    return jnp.sum(ys)


t0 = True
t1 = True
ts = None
y0__args__term = y0_args_term0

saveat = diffrax.SaveAt(t0=t0, t1=t1, ts=ts)

inexact, static = eqx.partition(y0__args__term, eqx.is_inexact_array)

def _run_inexact(inexact, saveat_, adjoint_):
    return _run(eqx.combine(inexact, static), saveat_, adjoint_)

_run_grad = eqx.filter_jit(jax.grad(_run_inexact))
_run_fwd_grad = eqx.filter_jit(jax.jacfwd(_run_inexact))
recursive_grads = _run_grad(inexact, saveat, diffrax.RecursiveCheckpointAdjoint())
forward_grads = _run_fwd_grad(inexact, saveat, diffrax.ForwardMode())
direct_grads = _run_grad(inexact, saveat, diffrax.DirectAdjoint())
backsolve_grads = _run_grad(
    inexact, saveat, diffrax.BacksolveAdjoint()
)

Forward errors with TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.

and

Backsolve errors with ValueError: Terms are not compatible with solver!

(this is from a fork of main).

I know there's this remark on closure with Backsolve (https://docs.kidger.site/diffrax/further_details/faq/#im-getting-a-customvjpexception), but back solve here has a different failure and forward I should be passing args, y0, terms directly into the function for forward in case it was the same.

Reproduce the supplied _run example with ForwardMode and BacksolveAdjoint, then compare the failures with the FAQ note on BacksolveAdjoint custom VJP behavior. Trace the adjoint and solver compatibility checks named by the errors; done means establishing whether these failures are expected and documenting or correcting the affected behavior.