From 8493548690d925ed44c0552030cb988be7c70afa Mon Sep 17 00:00:00 2001
From: Sam McCallum <sgmccallum@icloud.com>
Date: Sun, 9 Mar 2025 11:28:05 +0000
Subject: [PATCH] fix h calculation

---
 diffrax/_solver/srk.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/diffrax/_solver/srk.py b/diffrax/_solver/srk.py
index 56be17ba..a7d1e905 100644
--- a/diffrax/_solver/srk.py
+++ b/diffrax/_solver/srk.py
@@ -352,10 +352,7 @@ def step(
         else:
             ignore_stage_g = jnp.array(self.tableau.ignore_stage_g)
 
-        # time increment
-        h = t1 - t0
-
-        # First the drift related stuff
+        # # First the drift related stuff
         a = self._embed_a_lower(self.tableau.a, dtype)
         c = jnp.asarray(
             np.insert(self.tableau.c, 0, 0.0), dtype=complex_to_real_dtype(dtype)
@@ -380,6 +377,10 @@ def make_zeros_aux(leaf):
         # Now the diffusion related stuff
         # Brownian increment (and space-time Lévy area)
         bm_inc = diffusion.contr(t0, t1, use_levy=True)
+
+        # time increment
+        h = bm_inc.dt
+
         if not isinstance(bm_inc, self.minimal_levy_area):
             raise ValueError(
                 f"The Brownian increment {bm_inc} does not have the "