Merge branch 'dev'

Author: lockwo

.pre-commit-config.yaml

@@ -1,14 +1,14 @@
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.7
+    rev: v0.2.2
     hooks:
       - id: ruff  # linter
         types_or: [ python, pyi, jupyter ]
         args: [ --fix ]
       - id: ruff-format  # formatter
         types_or: [ python, pyi, jupyter ]
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.316
+    rev: v1.1.350
     hooks:
     - id: pyright
       additional_dependencies: [equinox, jax, jaxtyping, optax, optimistix, lineax, pytest, typing_extensions]

benchmarks/small_neural_ode.py

@@ -20,7 +20,7 @@
 class FuncTorch(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.func = torch.jit.script(  # pyright: ignore
+        self.func = torch.jit.script(
             torch.nn.Sequential(
                 torch.nn.Linear(4, 32),
                 torch.nn.Softplus(),
@@ -30,7 +30,7 @@ def __init__(self):
         )
 
     def forward(self, t, y):
-        return self.func(y)  # pyright: ignore
+        return self.func(y)
 
 
 class FuncJax(eqx.Module):
@@ -177,10 +177,10 @@ def run(multiple, grad, batch_size=64, t1=100):
     with torch.no_grad():
         func_jax = neural_ode_diffrax.func.func
         func_torch = neural_ode_torch.func.func
-        func_torch[0].weight.copy_(torch.tensor(np.asarray(func_jax.layers[0].weight)))  # pyright: ignore
-        func_torch[0].bias.copy_(torch.tensor(np.asarray(func_jax.layers[0].bias)))  # pyright: ignore
-        func_torch[2].weight.copy_(torch.tensor(np.asarray(func_jax.layers[1].weight)))  # pyright: ignore
-        func_torch[2].bias.copy_(torch.tensor(np.asarray(func_jax.layers[1].bias)))  # pyright: ignore
+        func_torch[0].weight.copy_(torch.tensor(np.asarray(func_jax.layers[0].weight)))
+        func_torch[0].bias.copy_(torch.tensor(np.asarray(func_jax.layers[0].bias)))
+        func_torch[2].weight.copy_(torch.tensor(np.asarray(func_jax.layers[1].weight)))
+        func_torch[2].bias.copy_(torch.tensor(np.asarray(func_jax.layers[1].bias)))
 
     y0_jax = jr.normal(jr.PRNGKey(1), (batch_size, 4))
     y0_torch = torch.tensor(np.asarray(y0_jax))

diffrax/__init__.py

@@ -13,7 +13,14 @@
     UnsafeBrownianPath as UnsafeBrownianPath,
     VirtualBrownianTree as VirtualBrownianTree,
 )
-from ._custom_types import LevyVal as LevyVal
+from ._custom_types import (
+    AbstractBrownianIncrement as AbstractBrownianIncrement,
+    AbstractSpaceTimeLevyArea as AbstractSpaceTimeLevyArea,
+    AbstractSpaceTimeTimeLevyArea as AbstractSpaceTimeTimeLevyArea,
+    BrownianIncrement as BrownianIncrement,
+    SpaceTimeLevyArea as SpaceTimeLevyArea,
+    SpaceTimeTimeLevyArea as SpaceTimeTimeLevyArea,
+)
 from ._event import (
     AbstractDiscreteTerminatingEvent as AbstractDiscreteTerminatingEvent,
     DiscreteTerminatingEvent as DiscreteTerminatingEvent,
@@ -37,6 +44,12 @@
 )
 from ._misc import adjoint_rms_seminorm as adjoint_rms_seminorm
 from ._path import AbstractPath as AbstractPath
+from ._progress_meter import (
+    AbstractProgressMeter as AbstractProgressMeter,
+    NoProgressMeter as NoProgressMeter,
+    TextProgressMeter as TextProgressMeter,
+    TqdmProgressMeter as TqdmProgressMeter,
+)
 from ._root_finder import (
     VeryChord as VeryChord,
     with_stepsize_controller_tols as with_stepsize_controller_tols,
@@ -59,6 +72,7 @@
     AbstractRungeKutta as AbstractRungeKutta,
     AbstractSDIRK as AbstractSDIRK,
     AbstractSolver as AbstractSolver,
+    AbstractSRK as AbstractSRK,
     AbstractStratonovichSolver as AbstractStratonovichSolver,
     AbstractWrappedSolver as AbstractWrappedSolver,
     Bosh3 as Bosh3,
@@ -68,6 +82,7 @@
     Dopri8 as Dopri8,
     Euler as Euler,
     EulerHeun as EulerHeun,
+    GeneralShARK as GeneralShARK,
     HalfSolver as HalfSolver,
     Heun as Heun,
     ImplicitEuler as ImplicitEuler,
@@ -84,8 +99,14 @@
     MultiButcherTableau as MultiButcherTableau,
     Ralston as Ralston,
     ReversibleHeun as ReversibleHeun,
+    SEA as SEA,
     SemiImplicitEuler as SemiImplicitEuler,
+    ShARK as ShARK,
     Sil3 as Sil3,
+    SlowRK as SlowRK,
+    SPaRK as SPaRK,
+    SRA1 as SRA1,
+    StochasticButcherTableau as StochasticButcherTableau,
     StratonovichMilstein as StratonovichMilstein,
     Tsit5 as Tsit5,
 )

diffrax/_adjoint.py

@@ -1,8 +1,8 @@
 import abc
 import functools as ft
 import warnings
-from collections.abc import Iterable
-from typing import Any, Optional, Union
+from collections.abc import Callable, Iterable
+from typing import Any, cast, Optional, Union
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -20,6 +20,9 @@
 from ._term import AbstractTerm, AdjointTerm
 
 
+ω = cast(Callable, ω)
+
+
 def _is_none(x):
     return x is None
 
@@ -128,6 +131,7 @@ def loop(
         init_state,
         passed_solver_state,
         passed_controller_state,
+        progress_meter,
     ) -> Any:
         """Runs the main solve loop. Subclasses can override this to provide custom
         backpropagation behaviour; see for example the implementation of
@@ -559,6 +563,7 @@ def _loop_backsolve_bwd(
     max_steps,
     throw,
     init_state,
+    progress_meter,
 ):
     assert discrete_terminating_event is None
 
@@ -567,7 +572,7 @@ def _loop_backsolve_bwd(
     # using them later.
     #
 
-    del perturbed, init_state, t1
+    del perturbed, init_state, t1, progress_meter
     ts, ys = residuals
     del residuals
     grad_final_state, _ = grad_final_state__aux_stats

diffrax/_autocitation.py

@@ -16,14 +16,23 @@
 from ._saveat import SubSaveAt
 from ._solver import (
     AbstractImplicitSolver,
+    AbstractItoSolver,
+    AbstractSRK,
+    AbstractStratonovichSolver,
     Dopri5,
     Dopri8,
+    GeneralShARK,
     Kvaerno3,
     Kvaerno4,
     Kvaerno5,
     LeapfrogMidpoint,
     ReversibleHeun,
+    SEA,
     SemiImplicitEuler,
+    ShARK,
+    SlowRK,
+    SPaRK,
+    SRA1,
     Tsit5,
 )
 from ._step_size_controller import PIDController
@@ -374,7 +383,15 @@ def _backsolve_rms_norm(adjoint):
 
 @citation_rules.append
 def _explicit_solver(solver, terms=None):
-    if not isinstance(solver, AbstractImplicitSolver) and not is_sde(terms):
+    if not isinstance(
+        solver,
+        (
+            AbstractImplicitSolver,
+            AbstractSRK,
+            AbstractItoSolver,
+            AbstractStratonovichSolver,
+        ),
+    ) and not is_sde(terms):
         return r"""
 % You are using an explicit solver, and may wish to cite the standard textbook:
 @book{hairer2008solving-i,
@@ -467,6 +484,12 @@ def _solvers(solver, saveat=None):
         Kvaerno5,
         ReversibleHeun,
         LeapfrogMidpoint,
+        ShARK,
+        SRA1,
+        SlowRK,
+        GeneralShARK,
+        SPaRK,
+        SEA,
     ):
         return (
             r"""

diffrax/_brownian/base.py

@@ -1,17 +1,25 @@
 import abc
-from typing import Optional, Union
+from typing import Optional, TypeVar, Union
 
 from equinox.internal import AbstractVar
 from jaxtyping import Array, PyTree
 
-from .._custom_types import LevyArea, LevyVal, RealScalarLike
+from .._custom_types import (
+    AbstractBrownianIncrement,
+    BrownianIncrement,
+    RealScalarLike,
+    SpaceTimeLevyArea,
+)
 from .._path import AbstractPath
 
 
-class AbstractBrownianPath(AbstractPath):
+_Control = TypeVar("_Control", bound=Union[PyTree[Array], AbstractBrownianIncrement])
+
+
+class AbstractBrownianPath(AbstractPath[_Control]):
     """Abstract base class for all Brownian paths."""
 
-    levy_area: AbstractVar[LevyArea]
+    levy_area: AbstractVar[type[Union[BrownianIncrement, SpaceTimeLevyArea]]]
 
     @abc.abstractmethod
     def evaluate(
@@ -20,7 +28,7 @@ def evaluate(
         t1: Optional[RealScalarLike] = None,
         left: bool = True,
         use_levy: bool = False,
-    ) -> Union[PyTree[Array], LevyVal]:
+    ) -> _Control:
         r"""Samples a Brownian increment $w(t_1) - w(t_0)$.
 
         Each increment has distribution $\mathcal{N}(0, t_1 - t_0)$.

diffrax/_brownian/path.py

@@ -10,7 +10,13 @@
 import lineax.internal as lxi
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
-from .._custom_types import levy_tree_transpose, LevyArea, LevyVal, RealScalarLike
+from .._custom_types import (
+    AbstractBrownianIncrement,
+    BrownianIncrement,
+    levy_tree_transpose,
+    RealScalarLike,
+    SpaceTimeLevyArea,
+)
 from .._misc import (
     force_bitcast_convert_type,
     is_tuple_of_ints,
@@ -42,25 +48,23 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: LevyArea = eqx.field(static=True)
+    levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]] = eqx.field(
+        static=True
+    )
     key: PRNGKeyArray
 
     def __init__(
         self,
         shape: Union[tuple[int, ...], PyTree[jax.ShapeDtypeStruct]],
         key: PRNGKeyArray,
-        levy_area: LevyArea = "",
+        levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]],
     ):
         self.shape = (
             jax.ShapeDtypeStruct(shape, lxi.default_floating_dtype())
             if is_tuple_of_ints(shape)
             else shape
         )
         self.key = key
-        if levy_area not in ["", "space-time"]:
-            raise ValueError(
-                f"levy_area must be one of '', 'space-time', but got {levy_area}."
-            )
         self.levy_area = levy_area
 
         if any(
@@ -84,7 +88,7 @@ def evaluate(
         t1: Optional[RealScalarLike] = None,
         left: bool = True,
         use_levy: bool = False,
-    ) -> Union[PyTree[Array], LevyVal]:
+    ) -> Union[PyTree[Array], AbstractBrownianIncrement]:
         del left
         if t1 is None:
             dtype = jnp.result_type(t0)
@@ -111,8 +115,8 @@ def evaluate(
             self.shape,
         )
         if use_levy:
-            out = levy_tree_transpose(self.shape, self.levy_area, out)
-            assert isinstance(out, LevyVal)
+            out = levy_tree_transpose(self.shape, out)
+            assert isinstance(out, (BrownianIncrement, SpaceTimeLevyArea))
         return out
 
     @staticmethod
@@ -121,25 +125,26 @@ def _evaluate_leaf(
         t1: RealScalarLike,
         key,
         shape: jax.ShapeDtypeStruct,
-        levy_area: str,
+        levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]],
         use_levy: bool,
     ):
         w_std = jnp.sqrt(t1 - t0).astype(shape.dtype)
+        w = jr.normal(key, shape.shape, shape.dtype) * w_std
+        dt = jnp.asarray(t1 - t0, dtype=shape.dtype)
 
-        if levy_area == "space-time":
+        if levy_area is SpaceTimeLevyArea:
             key, key_hh = jr.split(key, 2)
             hh_std = w_std / math.sqrt(12)
             hh = jr.normal(key_hh, shape.shape, shape.dtype) * hh_std
-        elif levy_area == "":
-            hh = None
+            levy_val = SpaceTimeLevyArea(dt=dt, W=w, H=hh)
+        elif levy_area is BrownianIncrement:
+            levy_val = BrownianIncrement(dt=dt, W=w)
         else:
             assert False
-        w = jr.normal(key, shape.shape, shape.dtype) * w_std
 
         if use_levy:
-            return LevyVal(dt=t1 - t0, W=w, H=hh, bar_H=None, K=None, bar_K=None)
-        else:
-            return w
+            return levy_val
+        return w
 
 
 UnsafeBrownianPath.__init__.__doc__ = """