fix test

Author: lockwo

diffrax/_integrate.py

@@ -760,8 +760,6 @@ def _promote(yi):
             _dtype = jnp.result_type(yi, time_dtype)  # noqa: F821
         return jnp.asarray(yi, dtype=_dtype)
 
-    if isinstance(solver, KLSolver):
-        y0 = (y0, 0.0)
     y0 = jtu.tree_map(_promote, y0)
     del timelikes
 
@@ -810,6 +808,9 @@ def _promote(yi):
                 "`UnsafeBrownianPath` cannot be used with adaptive step sizes."
             )
 
+    if isinstance(solver, KLSolver):
+        y0 = (y0, 0.0)
+        y0 = jtu.tree_map(_promote, y0)
     # Normalises time: if t0 > t1 then flip things around.
     direction = jnp.where(t0 < t1, 1, -1)
     t0 = t0 * direction

diffrax/_solver/kl.py

@@ -104,7 +104,7 @@ def vf(self, t: RealScalarLike, y: Y, args: Args) -> Tuple[VF, RealScalarLike]:
             self.drift1, self.drift2, self.diffusion, t, y, args, self.linear_solver
         )
 
-    def contr(self, t0: RealScalarLike, t1: RealScalarLike) -> Control:
+    def contr(self, t0: RealScalarLike, t1: RealScalarLike, **kwargs) -> Control:
         return t1 - t0
 
     def prod(self, vf: VF, control: RealScalarLike) -> Y:
@@ -120,7 +120,7 @@ def vf(self, t: RealScalarLike, y: Y, args: Args) -> Tuple[VF, RealScalarLike]:
         return vf, 0.0
 
     def contr(
-        self, t0: RealScalarLike, t1: RealScalarLike
+        self, t0: RealScalarLike, t1: RealScalarLike, **kwargs
     ) -> Tuple[Control, RealScalarLike]:
         return self.control_term.contr(t0, t1), 0.0