Custom VJP error with Custom PyTrees and RecursiveCheckpoint

To my understanding, the following MWE should be able to find gradients with respect to the args variable of the vector field function

from jax import numpy as jnp
import diffrax as dfx
import equinox as eqx
from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class CustomPytree:
  def __init__(self, p):
    self.p = p
    
  def tree_flatten(self):
    return (self.p,), None

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    return cls(*children)
  
  def __mul__(self, other):
    return CustomPytree(self.p * other.p)
  
u = CustomPytree(jnp.ones([32,32]))
c = CustomPytree(jnp.ones([32,32]))

saveat = dfx.SaveAt(t1=True)
solver = dfx.Tsit5()

vector_field = lambda t, x, args: args*x

@eqx.filter_jit
@eqx.filter_value_and_grad
def lossfn(args, u0):
  term = dfx.ODETerm(vector_field)
  
  sol = dfx.diffeqsolve(
    term, solver, 0, 10, 0.1, u0, 
    saveat=saveat, 
    args=args,
    adjoint=dfx.RecursiveCheckpointAdjoint()
  )

  return jnp.mean(sol.ys.p)**2
  
outs = lossfn(c, u)

But I get a TypeError: Custom VJP ... etc.
Taking the gradient against the CustomPytree for generic jax function works, for example this works fine

f = lambda c, u: jnp.mean((c*u).p)**2
val, grad = eqx.filter_value_and_grad(f)(c, u)

Am I missing something obvious?

Thanks in advance!

EDIT: I sent this before finishing to write it, sorry 😄

Run the supplied MWE starting at lossfn, the CustomPytree definition, and RecursiveCheckpointAdjoint to reproduce the Custom VJP TypeError. Compare this with the working generic JAX-function example and trace the gradient path through diffeqsolve; done means gradients with respect to args work for the custom PyTree case without the reported error.