some fixes

Author: lockwo

diffrax/_integrate.py

@@ -340,7 +340,6 @@ def cond_fun(state):
 
     def body_fun_aux(state):
         state = _handle_static(state)
-
         #
         # Actually do some differential equation solving! Make numerical steps, adapt
         # step sizes, all that jazz.
@@ -1105,7 +1104,11 @@ def _promote(yi):
             terms = MultiTerm(*terms)
 
     if path_state is None:
-        path_state = terms.init(t0, t1, y0, args)
+        path_state = jax.tree.map(
+            lambda term: term.init(t0, t1, y0, args),
+            terms,
+            is_leaf=lambda x: isinstance(x, AbstractTerm),
+        )
 
     # Error checking for term compatibility
     _assert_term_compatible(
@@ -1252,7 +1255,11 @@ def _subsaveat_direction_fn(x):
 
     if path_state is None:
         passed_path_state = False
-        path_state = terms.init(t0, tnext, y0, args)
+        path_state = jax.tree.map(
+            lambda term: term.init(t0, tnext, y0, args),
+            terms,
+            is_leaf=lambda x: isinstance(x, AbstractTerm),
+        )
     else:
         passed_path_state = True
 

diffrax/_solver/foster_langevin_srk.py

@@ -260,6 +260,7 @@ def init(
         evaluation of grad_f.
         """
         drift, diffusion = terms.terms
+        drift_path, diffusion_path = path_state
         (
             gamma_drift,
             u_drift,
@@ -271,7 +272,7 @@ def init(
         # is this the only solver class that has `init` depend on the path state?
         # feels irksome to change everything for one class, but I'm going to make
         # `init` now depend on path state for the sake of generality
-        h, _ = drift.contr(t0, t1, path_state)
+        h, _ = drift.contr(t0, t1, drift_path)
         x0, v0 = y0
 
         gamma = broadcast_underdamped_langevin_arg(gamma_drift, x0, "gamma")
@@ -390,7 +391,6 @@ def step(
         drift, diffusion = terms.terms
         drift_path, diffusion_path = path_state
 
-
         h, drift_path = drift.contr(t0, t1, drift_path)
         h_prev = st.h
         tay: PyTree[_Coeffs] = st.taylor_coeffs

diffrax/_solver/milstein.py

@@ -69,6 +69,8 @@ def init(
     ) -> _SolverState:
         return None
 
+    # TODO, a bunch of these solvers have tuple requirements, we can type the
+    # _PathState to be the same pytree.
     def step(
         self,
         terms: MultiTerm[
@@ -84,9 +86,10 @@ def step(
     ) -> tuple[Y, _ErrorEstimate, DenseInfo, _SolverState, _PathState, RESULTS]:
         del solver_state, made_jump
         drift, diffusion = terms.terms
-        # should these be same path state?
-        dt, _ = 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_prod = drift.vf_prod(t0, y0, args, dt)
         g0_prod = diffusion.vf_prod(t0, y0, args, dw)
@@ -98,7 +101,14 @@ def _to_jvp(_y0):
         y1 = (y0**ω + f0_prod**ω + g0_prod**ω + 0.5 * v0_prod**ω).ω
 
         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,
@@ -167,8 +177,9 @@ def step(
     ) -> tuple[Y, _ErrorEstimate, DenseInfo, _SolverState, _PathState, RESULTS]:
         del solver_state, made_jump
         drift, diffusion = terms.terms
-        Δt, path_state = drift.contr(t0, t1, path_state)
-        Δw, path_state = diffusion.contr(t0, t1, path_state)
+        drift_path, diffusion_path = path_state
+        Δt, drift_path = drift.contr(t0, t1, drift_path)
+        Δw, diffusion_path = diffusion.contr(t0, t1, diffusion_path)
 
         #
         # So this is a bit involved, largely because of the generality that the rest of
@@ -379,7 +390,14 @@ def _dot(_, _v0):
         #
 
         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/_solver/semi_implicit_euler.py

@@ -62,16 +62,24 @@ def step(
         del solver_state, made_jump
 
         term_1, term_2 = terms
+        path_state1, path_state2 = path_state
         y0_1, y0_2 = y0
 
-        control1, path_state = term_1.contr(t0, t1, path_state)
-        control2, path_state = term_2.contr(t0, t1, path_state)
+        control1, path_state1 = term_1.contr(t0, t1, path_state1)
+        control2, path_state2 = term_2.contr(t0, t1, path_state2)
         y1_1 = (y0_1**ω + term_1.vf_prod(t0, y0_2, args, control1) ** ω).ω
         y1_2 = (y0_2**ω + term_2.vf_prod(t0, y1_1, args, control2) ** ω).ω
 
         y1 = (y1_1, y1_2)
         dense_info = dict(y0=y0, y1=y1)
-        return y1, None, dense_info, None, path_state, RESULTS.successful
+        return (
+            y1,
+            None,
+            dense_info,
+            None,
+            (path_state1, path_state2),
+            RESULTS.successful,
+        )
 
     def func(
         self,

diffrax/_solver/srk.py

@@ -663,7 +663,14 @@ def compute_and_insert_kg_j(_w_kgs_in, _levylist_kgs_in):
 
         y1 = (y0**ω + drift_result**ω + diffusion_result**ω).ω
         dense_info = dict(y0=y0, y1=y1)
-        return y1, error, dense_info, None, (drift_path, diffusion_path), RESULTS.successful
+        return (
+            y1,
+            error,
+            dense_info,
+            None,
+            (drift_path, diffusion_path),
+            RESULTS.successful,
+        )
 
     def func(
         self,

diffrax/_term.py

@@ -675,7 +675,6 @@ def contr(
         control_state: PyTree,
         **kwargs,
     ) -> tuple[tuple[PyTree[ArrayLike], ...], tuple[PyTree, ...]]:
-
         contrs = [
             term.contr(t0, t1, state, **kwargs)
             for term, state in zip(self.terms, control_state)

test/test_solver.py

@@ -205,6 +205,9 @@ def test_everything_pytree(implicit, vf_expensive, adaptive):
     class Term(diffrax.AbstractTerm):
         coeff: float
 
+        def init(self, t0, t1, y0, args):
+            return None
+
         def vf(self, t, y, args):
             return {"f": -self.coeff * y["y"]}
 

test/test_underdamped_langevin.py

@@ -91,7 +91,6 @@ def test_shape(solver, dtype):
     sde = get_pytree_uld(t0, t1, dtype)
     bm = sde.get_bm(jr.key(5678), diffrax.SpaceTimeTimeLevyArea, tol=0.2)
     terms = sde.get_terms(bm)
-    print(terms)
 
     sol = diffeqsolve(
         terms, solver, t0, t1, dt0=dt0, y0=sde.y0, args=None, saveat=saveat