fixes

Author: lockwo

diffrax/_adjoint.py

@@ -909,9 +909,10 @@ def loop(
         throw,
         passed_solver_state,
         passed_controller_state,
+        passed_path_state,
         **kwargs,
     ):
-        del throw, passed_solver_state, passed_controller_state
+        del throw, passed_solver_state, passed_controller_state, passed_path_state
         inner_while_loop = eqx.Partial(_inner_loop, kind="lax")
         outer_while_loop = eqx.Partial(_outer_loop, kind="lax")
         # Support forward-mode autodiff.

diffrax/_integrate.py

@@ -181,6 +181,7 @@ def _check(term_cls, term, term_contr_kwargs, yi):
                 if not vf_type_compatible:
                     raise ValueError(f"Vector field term {term} is incompatible.")
 
+                term_contr_kwargs["control_state"] = term.init(0.0, 0.0, y, args)
                 contr = ft.partial(term.contr, **term_contr_kwargs)
                 # Work around https://github.com/google/jax/issues/21825
                 try:
@@ -387,7 +388,7 @@ def body_fun_aux(state):
 
         tprev = jnp.minimum(tprev, t1)
         tnext = _clip_to_end(tprev, tnext, t1, keep_step)
-        
+
         progress_meter_state = progress_meter.step(
             state.progress_meter_state, linear_rescale(t0, tprev, t1)
         )
@@ -1111,17 +1112,26 @@ def _promote(yi):
         )
 
     # Error checking for term compatibility
+
+    # try:
+    #     contr_kwargs = jtu.tree_map(
+    #         lambda _, x, y: jtu.tree_map(
+    #             lambda a, b: a | {"control_state": b},
+    #             x,
+    #             y,
+    #             is_leaf=lambda v: isinstance(v, dict),
+    #         ),
+    #         solver.term_structure,
+    #         solver.term_compatible_contr_kwargs,
+    #         path_state,
+    #         is_leaf=lambda z: isinstance(z, AbstractTerm)
+    #         and not isinstance(z, MultiTerm),
+    #     )
+    # except Exception as e:
+    #     raise ValueError("Terms are not compatible with solver!") from e
+
     _assert_term_compatible(
-        y0,
-        args,
-        terms,
-        solver.term_structure,
-        jtu.tree_map(
-            lambda x, y: x | {"control_state": y},
-            solver.term_compatible_contr_kwargs,
-            path_state,
-            is_leaf=lambda x: isinstance(x, dict),
-        ),
+        y0, args, terms, solver.term_structure, solver.term_compatible_contr_kwargs
     )
 
     if is_sde(terms):

diffrax/_solver/euler_heun.py

@@ -68,8 +68,9 @@ def step(
         del solver_state, made_jump
 
         drift, diffusion = terms.terms
-        dt, path_state = drift.contr(t0, t1, path_state)
-        dW, path_state = diffusion.contr(t0, t1, path_state)
+        drift_path, diffusion_path = path_state
+        dt, drift_path = drift.contr(t0, t1, drift_path)
+        dW, diffusion_path = diffusion.contr(t0, t1, diffusion_path)
 
         f0 = drift.vf_prod(t0, y0, args, dt)
         g0 = diffusion.vf_prod(t0, y0, args, dW)
@@ -80,7 +81,14 @@ def step(
         y1 = (y0**ω + f0**ω + 0.5 * (g0**ω + g_prime**ω)).ω
 
         dense_info = dict(y0=y0, y1=y1)
-        return y1, None, dense_info, None, path_state, RESULTS.successful
+        return (
+            y1,
+            None,
+            dense_info,
+            None,
+            (drift_path, diffusion_path),
+            RESULTS.successful,
+        )
 
     def func(
         self,

diffrax/_term.py

@@ -777,7 +777,9 @@ def is_vf_expensive(
         ],
         args: Args,
     ) -> bool:
-        control_struct = eqx.filter_eval_shape(self.contr, t0, t1)
+        control_struct = eqx.filter_eval_shape(
+            self.contr, t0, t1, self.term.init(t0, t1, y, args)
+        )
         if sum(c.size for c in jtu.tree_leaves(control_struct)) in (0, 1):
             return False
         else:

test/test_integrate.py

@@ -14,7 +14,7 @@
 import scipy.stats
 from diffrax import ControlTerm, MultiTerm, ODETerm
 from equinox.internal import ω
-from jaxtyping import Array, ArrayLike, Float
+from jaxtyping import Array, ArrayLike, Float, PyTree
 
 from .helpers import (
     all_ode_solvers,
@@ -638,6 +638,10 @@ class TestSolver(diffrax.Euler):
 
 
 def test_term_compatibility_pytree():
+    class _TestState(eqx.Module):
+        y: PyTree
+        state: PyTree
+
     class TestSolver(diffrax.AbstractSolver):
         term_structure = {
             "a": diffrax.ODETerm,
@@ -661,14 +665,20 @@ def init(self, terms, t0, t1, y0, args, path_state):
             return None
 
         def step(self, terms, t0, t1, y0, args, solver_state, made_jump, path_state):
-            def _step(_term, _y):
-                control = _term.contr(t0, t1)
-                return _y + _term.vf_prod(t0, _y, args, control)
+            def _step(_term, _y, state):
+                control, new_state = _term.contr(t0, t1, state)
+                return _TestState(_y + _term.vf_prod(t0, _y, args, control), new_state)
 
             _is_term = lambda x: isinstance(x, diffrax.AbstractTerm)
-            y1 = jtu.tree_map(_step, terms, y0, is_leaf=_is_term)
+            output = jtu.tree_map(_step, terms, y0, path_state, is_leaf=_is_term)
+            y1 = jtu.tree_map(
+                lambda x: x.y, output, is_leaf=lambda x: isinstance(x, _TestState)
+            )
+            path_state = jtu.tree_map(
+                lambda x: x.state, output, is_leaf=lambda x: isinstance(x, _TestState)
+            )
             dense_info = dict(y0=y0, y1=y1)
-            return y1, None, dense_info, None, None, diffrax.RESULTS.successful
+            return y1, None, dense_info, None, path_state, diffrax.RESULTS.successful
 
         def func(self, terms, t0, y0, args):
             assert False