fix

patrick-kidger · Jan 22, 2025 · d0f161c · d0f161c
1 parent 1ad8dad
commit d0f161c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/diffrax/_brownian/path.py b/diffrax/_brownian/path.py
@@ -119,9 +119,9 @@ def _generate_noise(
         max_steps: int,
     ) -> Float[Array, "..."]:
         if self.levy_area is SpaceTimeTimeLevyArea:
-            noise = jr.normal(key, (3, max_steps, *shape.shape), shape.dtype)
+            noise = jr.normal(key, (max_steps, 3, *shape.shape), shape.dtype)
         elif self.levy_area is SpaceTimeLevyArea:
-            noise = jr.normal(key, (2, max_steps, *shape.shape), shape.dtype)
+            noise = jr.normal(key, (max_steps, 2, *shape.shape), shape.dtype)
         elif self.levy_area is BrownianIncrement:
             noise = jr.normal(key, (max_steps, *shape.shape), shape.dtype)
         else:

diff --git a/test/test_brownian.py b/test/test_brownian.py
@@ -131,11 +131,13 @@ def test_statistics(ctr, levy_area, use_levy):
     def _eval(key):
         if ctr is diffrax.UnsafeBrownianPath:
             path = ctr(shape=(), key=key, levy_area=levy_area)
+            state = path.init(t0, t1, None, None)
         elif ctr is diffrax.VirtualBrownianTree:
             path = ctr(t0=0, t1=5, tol=2**-5, shape=(), key=key, levy_area=levy_area)
+            state = path.init(t0, t1, None, None)
         else:
             assert False
-        return path.evaluate(t0, t1, use_levy=use_levy)
+        return path(t0, state, t1, use_levy=use_levy)[0]
 
     values = jax.vmap(_eval)(keys)
     if use_levy: