tests + examples

Author: lockwo

benchmarks/stateful_paths.py

@@ -268,6 +268,7 @@ def step(y, dW):
 New UBP + Precompute: 0.002430
 Pure Jax: 0.002799
 
+(these are out of date)
 Results on A100 GPU:
 VBT: 3.881952
 Old UBP: 0.337173

diffrax/_adjoint.py

@@ -377,6 +377,9 @@ def loop(
         #         "Cannot reverse-mode autodifferentiate when using "
         #         "`UnsafeBrownianPath`."
         #     )
+        # if is_unsafe_sde(terms):
+        #     kind = "lax"
+        #     msg = None
         if max_steps is None:
             kind = "lax"
             msg = (
@@ -836,7 +839,10 @@ def loop(
             raise NotImplementedError(
                 "Cannot use `adjoint=BacksolveAdjoint()` with `saveat=SaveAt(fn=...)`."
             )
-        # is this still true with DirectAdjoint?
+        # is this still true with DirectBP?
+        # it seems to give inaccurate results, so not currently, but seems doable
+        # might just require more careful thinking about path state management
+        # and more knowledge about continuous adjoints than I have currently
         if is_unsafe_sde(terms):
             raise ValueError(
                 "`adjoint=BacksolveAdjoint()` does not support `UnsafeBrownianPath`. "

diffrax/_brownian/base.py

@@ -9,6 +9,7 @@
     BrownianIncrement,
     RealScalarLike,
     SpaceTimeLevyArea,
+    SpaceTimeTimeLevyArea
 )
 from .._path import AbstractPath
 
@@ -20,7 +21,7 @@
 class AbstractBrownianPath(AbstractPath[_Control, _BrownianState]):
     """Abstract base class for all Brownian paths."""
 
-    levy_area: AbstractVar[type[Union[BrownianIncrement, SpaceTimeLevyArea]]]
+    levy_area: AbstractVar[type[Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]]]
 
     @abc.abstractmethod
     def __call__(

diffrax/_brownian/path.py

@@ -17,6 +17,7 @@
     BrownianIncrement,
     levy_tree_transpose,
     RealScalarLike,
+    IntScalarLike,
     SpaceTimeLevyArea,
     SpaceTimeTimeLevyArea,
     Y,
@@ -31,7 +32,7 @@
 
 _Control = Union[PyTree[Array], AbstractBrownianIncrement]
 _BrownianState: TypeAlias = Union[
-    tuple[None, PyTree[Array], int], tuple[PRNGKeyArray, None, None]
+    tuple[None, PyTree[Array], IntScalarLike], tuple[PRNGKeyArray, None, None]
 ]
 
 
@@ -73,10 +74,10 @@ class DirectBrownianPath(AbstractBrownianPath[_Control, _BrownianState]):
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
+    key: PRNGKeyArray
     levy_area: type[
         Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
     ] = eqx.field(static=True)
-    key: PRNGKeyArray
     precompute: Optional[int] = eqx.field(static=True)
 
     def __init__(
@@ -116,7 +117,7 @@ def _generate_noise(
         key: PRNGKeyArray,
         shape: jax.ShapeDtypeStruct,
         max_steps: int,
-    ) -> Float[Array, "levy_dims shape"]:
+    ) -> Float[Array, "..."]:
         # TODO: merge into a single jr.normal call
         if self.levy_area is SpaceTimeTimeLevyArea:
             noise = jr.normal(key, (3, max_steps, *shape.shape), shape.dtype)
@@ -254,7 +255,7 @@ def _evaluate_leaf_precomputed(
             Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
         ],
         use_levy: bool,
-        noises: Float[Array, "levy_dims shape"],
+        noises: Float[Array, "..."],
     ):
         w_std = jnp.sqrt(t1 - t0).astype(shape.dtype)
         dt = jnp.asarray(t1 - t0, dtype=complex_to_real_dtype(shape.dtype))

examples/underdamped_langevin_example.ipynb

@@ -237,6 +237,18 @@ def test_direct_brownian():
         final_activation=jnp.tanh,
         key=diffusionkey,
     )
+    class Field(eqx.Module):
+        force: eqx.nn.MLP
+
+        def __call__(self, t, y, args):
+            return self.force(y)
+
+    class DiffusionField(eqx.Module):
+        force: eqx.nn.MLP
+
+        def __call__(self, t, y, args):
+            return lx.DiagonalLinearOperator(self.force(y))
+
     y0 = jr.normal(ykey, (3,))
 
     k1, k2, k3 = jax.random.split(key, 3)
@@ -250,25 +262,25 @@ def test_direct_brownian():
     )
 
     vbt_terms = diffrax.MultiTerm(
-        diffrax.ODETerm(lambda t, y, args: drift_mlp(y)),
+        diffrax.ODETerm(Field(drift_mlp)),
         diffrax.ControlTerm(
-            lambda t, y, args: lx.DiagonalLinearOperator(diffusion_mlp(y)), vbt
+            DiffusionField(diffusion_mlp), vbt
         ),
     )
     dbp_terms = diffrax.MultiTerm(
-        diffrax.ODETerm(lambda t, y, args: drift_mlp(y)),
+        diffrax.ODETerm(Field(drift_mlp)),
         diffrax.ControlTerm(
-            lambda t, y, args: lx.DiagonalLinearOperator(diffusion_mlp(y)), dbp
+            DiffusionField(diffusion_mlp), dbp
         ),
     )
     dbp_pre_terms = diffrax.MultiTerm(
-        diffrax.ODETerm(lambda t, y, args: drift_mlp(y)),
+        diffrax.ODETerm(Field(drift_mlp)),
         diffrax.ControlTerm(
-            lambda t, y, args: lx.DiagonalLinearOperator(diffusion_mlp(y)), dbp_pre
+            DiffusionField(diffusion_mlp), dbp_pre
         ),
     )
 
-    solver = diffrax.GeneralShARK()
+    solver = diffrax.Heun()
 
     y0_args_term0 = (y0, None, vbt_terms)
     y0_args_term1 = (y0, None, dbp_terms)
@@ -307,7 +319,7 @@ def _run_finite_diff(y0__args__term, saveat, adjoint):
     for t0 in (True, False):
         for t1 in (True, False):
             for ts in (None, [0.3], [2.0], [9.5], [1.0, 7.0], [0.3, 7.0, 9.5]):
-                for y0__args__term in (y0_args_term0,):#, y0_args_term1, y0_args_term2):
+                for i, y0__args__term in enumerate((y0_args_term0, y0_args_term1, y0_args_term2)):
                     if t0 is False and t1 is False and ts is None:
                         continue
 
@@ -329,17 +341,20 @@ def _run_inexact(inexact, saveat, adjoint):
                     recursive_grads = _run_grad(
                         inexact, saveat, diffrax.RecursiveCheckpointAdjoint()
                     )
-                    # backsolve_grads = _run_grad(
-                    #     inexact, saveat, diffrax.BacksolveAdjoint()
-                    # )
+                    if i == 0:
+                        backsolve_grads = _run_grad(
+                            inexact, saveat, diffrax.BacksolveAdjoint()
+                        )
+                        assert tree_allclose(fd_grads, backsolve_grads[0], atol=1e-3)
+                        
                     forward_grads = _run_fwd_grad(
                         inexact, saveat, diffrax.ForwardMode()
                     )
+                    # TODO: fix via https://github.com/patrick-kidger/equinox/issues/923
                     # direct_grads = _run_grad(inexact, saveat, diffrax.DirectAdjoint())
-                    # assert tree_allclose(fd_grads, direct_grads[0])
-                    assert tree_allclose(fd_grads, recursive_grads, atol=1e-5)
-                    # assert tree_allclose(fd_grads, backsolve_grads, atol=1e-5)
-                    assert tree_allclose(fd_grads, forward_grads, atol=1e-5)
+                    # assert tree_allclose(fd_grads, direct_grads[0], atol=1e-3)
+                    assert tree_allclose(fd_grads, recursive_grads[0], atol=1e-3)
+                    assert tree_allclose(fd_grads, forward_grads[0], atol=1e-3)
 
 
 def test_adjoint_seminorm():