Parametric control types

.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]

diffrax/__init__.py

@@ -13,7 +13,11 @@
     UnsafeBrownianPath as UnsafeBrownianPath,
     VirtualBrownianTree as VirtualBrownianTree,
 )
-from ._custom_types import LevyVal as LevyVal
+from ._custom_types import (
+    AbstractBrownianReturn as AbstractBrownianReturn,
+    BrownianIncrement as BrownianIncrement,
+    SpaceTimeLevyArea as SpaceTimeLevyArea,
+)
 from ._event import (
     AbstractDiscreteTerminatingEvent as AbstractDiscreteTerminatingEvent,
     DiscreteTerminatingEvent as DiscreteTerminatingEvent,

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
 

diffrax/_brownian/base.py

@@ -1,17 +1,20 @@
 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 AbstractBrownianReturn, RealScalarLike
 from .._path import AbstractPath
 
 
-class AbstractBrownianPath(AbstractPath):
+_Control = TypeVar("_Control", bound=Union[PyTree[Array], AbstractBrownianReturn])
+
+
+class AbstractBrownianPath(AbstractPath[_Control]):
     """Abstract base class for all Brownian paths."""
 
-    levy_area: AbstractVar[LevyArea]
+    levy_area: AbstractVar[type[AbstractBrownianReturn]]
 
     @abc.abstractmethod
     def evaluate(
@@ -20,7 +23,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,12 @@
 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 (
+    BrownianIncrement,
+    levy_tree_transpose,
+    RealScalarLike,
+    SpaceTimeLevyArea,
+)
 from .._misc import (
     force_bitcast_convert_type,
     is_tuple_of_ints,
@@ -42,26 +47,24 @@ 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
+        self.levy_area = levy_area  # pyright: ignore[reportIncompatibleVariableOverride]
 
         if any(
             not jnp.issubdtype(x.dtype, jnp.inexact)
@@ -70,11 +73,11 @@ def __init__(
             raise ValueError("UnsafeBrownianPath dtypes all have to be floating-point.")
 
     @property
-    def t0(self):
+    def t0(self):  # pyright: ignore[reportIncompatibleVariableOverride]
         return -jnp.inf
 
     @property
-    def t1(self):
+    def t1(self):  # pyright: ignore[reportIncompatibleVariableOverride]
         return jnp.inf
 
     @eqx.filter_jit
@@ -84,7 +87,7 @@ def evaluate(
         t1: Optional[RealScalarLike] = None,
         left: bool = True,
         use_levy: bool = False,
-    ) -> Union[PyTree[Array], LevyVal]:
+    ) -> Union[PyTree[Array], BrownianIncrement, SpaceTimeLevyArea]:
         del left
         if t1 is None:
             dtype = jnp.result_type(t0)
@@ -112,7 +115,7 @@ def evaluate(
         )
         if use_levy:
             out = levy_tree_transpose(self.shape, out)
-            assert isinstance(out, LevyVal)
+            assert isinstance(out, (BrownianIncrement, SpaceTimeLevyArea))
         return out
 
     @staticmethod
@@ -121,25 +124,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 = t1 - t0
 
-        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, K=None)
+        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, K=None)
-        else:
-            return w
+            return levy_val
+        return w
 
 
 UnsafeBrownianPath.__init__.__doc__ = """