update jit results

patrick-kidger · Feb 9, 2025 · 20e700d · 20e700d
1 parent 7865a16
commit 20e700d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/benchmarks/stateful_paths.py b/benchmarks/stateful_paths.py
@@ -74,7 +74,6 @@ def __call__(
     ):
         return self.evaluate(t0, t1, left, use_levy), brownian_state
 
-    @eqx.filter_jit
     def evaluate(
         self,
         t0,
@@ -270,9 +269,9 @@ def step(y, dW):
 
 Results on A100 GPU:
 VBT: 2.275057
-Old UBP: 0.092015
-New UBP: 0.125904
-New UBP + Precompute: 0.108587
+Old UBP: 0.112461
+New UBP: 0.126370
+New UBP + Precompute: 0.111837
 Pure Jax: 0.261937
 
 For small ndt (e.g. 100) the pure jax is faster, but the diffrax overhead

diff --git a/diffrax/_brownian/path.py b/diffrax/_brownian/path.py
@@ -153,6 +153,7 @@ def init(
             key = self.key
             return key, noise, counter
 
+    @eqx.filter_jit
     def __call__(
         self,
         t0: RealScalarLike,