pyupgrade, sort pyproject.

Author: patrick-kidger

.pre-commit-config.yaml

@@ -1,4 +1,12 @@
 repos:
+  - repo: local
+    hooks:
+      - id: sort_pyproject
+        name: sort_pyproject
+        entry: toml-sort -i --sort-table-keys --sort-inline-tables
+        language: python
+        files: ^pyproject\.toml$
+        additional_dependencies: ["toml-sort==0.23.1"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.2.2
     hooks:

benchmarks/brownian_tree_times.py

@@ -7,8 +7,7 @@
 """
 
 import timeit
-from typing import cast, Optional, Union
-from typing_extensions import TypeAlias
+from typing import cast, TypeAlias
 
 import diffrax
 import equinox as eqx
@@ -24,7 +23,7 @@
 from jaxtyping import Array, Float, PRNGKeyArray, PyTree, Real
 
 
-RealScalarLike: TypeAlias = Real[Union[int, float, Array, np.ndarray], ""]
+RealScalarLike: TypeAlias = Real[int | float | Array | np.ndarray, ""]
 
 
 class _State(eqx.Module):
@@ -68,7 +67,7 @@ def levy_area(self):
     def evaluate(
         self,
         t0: RealScalarLike,
-        t1: Optional[RealScalarLike] = None,
+        t1: RealScalarLike | None = None,
         left: bool = True,
         use_levy: bool = False,
     ) -> PyTree[Array]:

benchmarks/old_pid_controller.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import cast, Optional, TypeVar
+from typing import cast, TypeVar
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -80,7 +80,7 @@ def intermediate(carry):
 
 
 _ControllerState = TypeVar("_ControllerState")
-_Dt0 = TypeVar("_Dt0", None, RealScalarLike, Optional[RealScalarLike])
+_Dt0 = TypeVar("_Dt0", None, RealScalarLike, RealScalarLike | None)
 
 _PidState = tuple[
     BoolScalarLike, BoolScalarLike, RealScalarLike, RealScalarLike, RealScalarLike
@@ -95,7 +95,7 @@ def _none_or_array(x):
 
 
 class OldPIDController(
-    AbstractAdaptiveStepSizeController[_PidState, Optional[RealScalarLike]]
+    AbstractAdaptiveStepSizeController[_PidState, RealScalarLike | None]
 ):
     r"""See the doc of diffrax.PIDController for more information."""
 
@@ -104,20 +104,20 @@ class OldPIDController(
     pcoeff: RealScalarLike = 0
     icoeff: RealScalarLike = 1
     dcoeff: RealScalarLike = 0
-    dtmin: Optional[RealScalarLike] = None
-    dtmax: Optional[RealScalarLike] = None
+    dtmin: RealScalarLike | None = None
+    dtmax: RealScalarLike | None = None
     force_dtmin: bool = True
-    step_ts: Optional[Real[Array, " steps"]] = eqx.field(
+    step_ts: Real[Array, " steps"] | None = eqx.field(
         default=None, converter=_none_or_array
     )
-    jump_ts: Optional[Real[Array, " jumps"]] = eqx.field(
+    jump_ts: Real[Array, " jumps"] | None = eqx.field(
         default=None, converter=_none_or_array
     )
     factormin: RealScalarLike = 0.2
     factormax: RealScalarLike = 10.0
     norm: Callable[[PyTree], RealScalarLike] = optx.rms_norm
     safety: RealScalarLike = 0.9
-    error_order: Optional[RealScalarLike] = None
+    error_order: RealScalarLike | None = None
 
     def __check_init__(self):
         if self.jump_ts is not None and not jnp.issubdtype(
@@ -143,10 +143,10 @@ def init(
         t0: RealScalarLike,
         t1: RealScalarLike,
         y0: Y,
-        dt0: Optional[RealScalarLike],
+        dt0: RealScalarLike | None,
         args: Args,
         func: Callable[[PyTree[AbstractTerm], RealScalarLike, Y, Args], VF],
-        error_order: Optional[RealScalarLike],
+        error_order: RealScalarLike | None,
     ) -> tuple[RealScalarLike, _PidState]:
         del t1
         if dt0 is None:
@@ -195,7 +195,7 @@ def adapt_step_size(
         y0: Y,
         y1_candidate: Y,
         args: Args,
-        y_error: Optional[Y],
+        y_error: Y | None,
         error_order: RealScalarLike,
         controller_state: _PidState,
     ) -> tuple[
@@ -330,7 +330,7 @@ def _scale(_y0, _y1_candidate, _y_error):
         )
         return keep_step, next_t0, next_t1, made_jump, controller_state, result
 
-    def _get_error_order(self, error_order: Optional[RealScalarLike]) -> RealScalarLike:
+    def _get_error_order(self, error_order: RealScalarLike | None) -> RealScalarLike:
         # Attribute takes priority, if the user knows the correct error order better
         # than our guess.
         error_order = error_order if self.error_order is None else self.error_order

diffrax/_adjoint.py

@@ -2,7 +2,7 @@
 import functools as ft
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Any, cast, Optional, Union
+from typing import Any, cast
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -259,7 +259,7 @@ class RecursiveCheckpointAdjoint(AbstractAdjoint):
         ```
     """
 
-    checkpoints: Optional[int] = None
+    checkpoints: int | None = None
 
     def loop(
         self,
@@ -453,7 +453,7 @@ def _solve(inputs):
     _solve = _solve.__wrapped__  # pyright: ignore[reportFunctionMemberAccess]
 
 
-def _frozenset(x: Union[object, Iterable[object]]) -> frozenset[object]:
+def _frozenset(x: object | Iterable[object]) -> frozenset[object]:
     try:
         iter_x = iter(x)  # pyright: ignore
     except TypeError:

diffrax/_autocitation.py

@@ -2,7 +2,6 @@
 import inspect
 import re
 from collections.abc import Callable
-from typing import Optional
 
 import jax
 import jax.core
@@ -130,7 +129,7 @@ def citation(*args, **kwargs):
 _diffeqsignature = inspect.signature(diffeqsolve)
 
 
-citation_rules: list[Callable[..., Optional[str]]] = []
+citation_rules: list[Callable[..., str | None]] = []
 
 
 _thesis_cite = r"""
@@ -158,14 +157,14 @@ def citation(*args, **kwargs):
 _reference_regex = re.compile(r"```bibtex([^`]*)```")
 
 
-@ft.lru_cache(maxsize=None)
+@ft.cache
 def _parse_reference(obj) -> str:
     references = _reference_regex.findall(obj.__doc__)
     [reference] = [inspect.cleandoc(ref) for ref in references]
     return reference
 
 
-@ft.lru_cache(maxsize=None)
+@ft.cache
 def _parse_reference_multi(obj) -> list[str]:
     references = _reference_regex.findall(obj.__doc__)
     return [inspect.cleandoc(ref) for ref in references]

diffrax/_brownian/base.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Optional, TypeVar, Union
+from typing import TypeVar
 
 from equinox.internal import AbstractVar
 from jaxtyping import Array, PyTree
@@ -13,19 +13,19 @@
 from .._path import AbstractPath
 
 
-_Control = TypeVar("_Control", bound=Union[PyTree[Array], AbstractBrownianIncrement])
+_Control = TypeVar("_Control", bound=PyTree[Array] | AbstractBrownianIncrement)
 
 
 class AbstractBrownianPath(AbstractPath[_Control]):
     """Abstract base class for all Brownian paths."""
 
-    levy_area: AbstractVar[type[Union[BrownianIncrement, SpaceTimeLevyArea]]]
+    levy_area: AbstractVar[type[BrownianIncrement | SpaceTimeLevyArea]]
 
     @abc.abstractmethod
     def evaluate(
         self,
         t0: RealScalarLike,
-        t1: Optional[RealScalarLike] = None,
+        t1: RealScalarLike | None = None,
         left: bool = True,
         use_levy: bool = False,
     ) -> _Control:

diffrax/_brownian/path.py

@@ -1,5 +1,5 @@
 import math
-from typing import cast, Optional, Union
+from typing import cast
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -63,16 +63,16 @@ class UnsafeBrownianPath(AbstractBrownianPath):
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
     levy_area: type[
-        Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+        BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
     ] = eqx.field(static=True)
     key: PRNGKeyArray
 
     def __init__(
         self,
-        shape: Union[tuple[int, ...], PyTree[jax.ShapeDtypeStruct]],
+        shape: tuple[int, ...] | PyTree[jax.ShapeDtypeStruct],
         key: PRNGKeyArray,
         levy_area: type[
-            Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+            BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
         ] = BrownianIncrement,
     ):
         """**Arguments:**
@@ -112,10 +112,10 @@ def t1(self):
     def evaluate(
         self,
         t0: RealScalarLike,
-        t1: Optional[RealScalarLike] = None,
+        t1: RealScalarLike | None = None,
         left: bool = True,
         use_levy: bool = False,
-    ) -> Union[PyTree[Array], AbstractBrownianIncrement]:
+    ) -> PyTree[Array] | AbstractBrownianIncrement:
         """Implements [`diffrax.AbstractBrownianPath.evaluate`][]."""
         del left
         if t1 is None:
@@ -153,9 +153,7 @@ def _evaluate_leaf(
         t1: RealScalarLike,
         key,
         shape: jax.ShapeDtypeStruct,
-        levy_area: type[
-            Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
-        ],
+        levy_area: type[BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea],
         use_levy: bool,
     ):
         w_std = jnp.sqrt(t1 - t0).astype(shape.dtype)

diffrax/_brownian/tree.py

@@ -1,6 +1,5 @@
 import math
-from typing import Literal, Optional, TypeVar, Union
-from typing_extensions import TypeAlias
+from typing import Literal, TypeAlias, TypeVar
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -70,10 +69,10 @@
 class _LevyVal(eqx.Module):
     dt: RealScalarLike
     W: Array
-    H: Optional[Array]
-    bar_H: Optional[Array]
-    K: Optional[Array]
-    bar_K: Optional[Array]
+    H: Array | None
+    bar_H: Array | None
+    K: Array | None
+    bar_K: Array | None
 
     def __check_init__(self):
         if self.H is None:
@@ -88,8 +87,8 @@ class _State(eqx.Module):
     s: RealScalarLike  # starting time of the interval
     w_s_u_su: FloatTriple  # W_s, W_u, W_{s,u}
     key: PRNGKeyArray
-    bhh_s_u_su: Optional[FloatTriple]  # \bar{H}_s, _u, _{s,u}
-    bkk_s_u_su: Optional[FloatTriple]  # \bar{K}_s, _u, _{s,u}
+    bhh_s_u_su: FloatTriple | None  # \bar{H}_s, _u, _{s,u}
+    bkk_s_u_su: FloatTriple | None  # \bar{K}_s, _u, _{s,u}
 
 
 def _levy_diff(_, x0: _LevyVal, x1: _LevyVal) -> AbstractBrownianIncrement:
@@ -161,8 +160,8 @@ def _make_levy_val(_, x: _LevyVal) -> AbstractBrownianIncrement:
 
 
 def _split_interval(
-    pred: BoolScalarLike, x_stu: Optional[FloatTriple], x_st_tu: Optional[FloatDouble]
-) -> Optional[FloatTriple]:
+    pred: BoolScalarLike, x_stu: FloatTriple | None, x_st_tu: FloatDouble | None
+) -> FloatTriple | None:
     if x_stu is None:
         assert x_st_tu is None
         return None
@@ -237,7 +236,7 @@ class VirtualBrownianTree(AbstractBrownianPath):
     tol: RealScalarLike
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
     levy_area: type[
-        Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+        BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
     ] = eqx.field(static=True)
     key: PyTree[PRNGKeyArray]
     _spline: _Spline = eqx.field(static=True)
@@ -248,10 +247,10 @@ def __init__(
         t0: RealScalarLike,
         t1: RealScalarLike,
         tol: RealScalarLike,
-        shape: Union[tuple[int, ...], PyTree[jax.ShapeDtypeStruct]],
+        shape: tuple[int, ...] | PyTree[jax.ShapeDtypeStruct],
         key: PRNGKeyArray,
         levy_area: type[
-            Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+            BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
         ] = BrownianIncrement,
         _spline: _Spline = "sqrt",
     ):
@@ -327,10 +326,10 @@ def is_dt(z):
     def evaluate(
         self,
         t0: RealScalarLike,
-        t1: Optional[RealScalarLike] = None,
+        t1: RealScalarLike | None = None,
         left: bool = True,
         use_levy: bool = False,
-    ) -> Union[PyTree[Array], AbstractBrownianIncrement]:
+    ) -> PyTree[Array] | AbstractBrownianIncrement:
         """Implements [`diffrax.AbstractBrownianPath.evaluate`][]."""
         del left
         t0 = eqxi.nondifferentiable(t0, name="t0")
@@ -532,9 +531,9 @@ def _body_fun(_state: _State):
                     f" 'zero' splines are permitted, got {self._spline}."
                 )
 
-            hat_w_sr, hat_hh_sr, hat_kk_sr = [
+            hat_w_sr, hat_hh_sr, hat_kk_sr = (
                 x.squeeze(axis=-1) for x in jnp.split(hat_y, 3, axis=-1)
-            ]
+            )
             assert hat_w_sr.shape == hat_hh_sr.shape == hat_kk_sr.shape == shape
 
             w_sr = w_mean + hat_w_sr
@@ -631,10 +630,10 @@ def _brownian_arch(
         FloatTriple,
         FloatDouble,
         tuple[PRNGKeyArray, PRNGKeyArray],
-        Optional[FloatTriple],
-        Optional[FloatDouble],
-        Optional[FloatTriple],
-        Optional[FloatDouble],
+        FloatTriple | None,
+        FloatDouble | None,
+        FloatTriple | None,
+        FloatDouble | None,
     ]:
         r"""For `t = (s+u)/2` evaluates `w_t` and (optionally) `bhh_t`
          conditioned on `w_s`, `w_u`, `bhh_s`, `bhh_u`, where

diffrax/_custom_types.py

@@ -1,4 +1,4 @@
-from typing import Any, TYPE_CHECKING, Union
+from typing import Any, TYPE_CHECKING
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -18,10 +18,10 @@
 
 
 if TYPE_CHECKING:
-    BoolScalarLike = Union[bool, Array, np.ndarray]
-    FloatScalarLike = Union[float, Array, np.ndarray]
-    IntScalarLike = Union[int, Array, np.ndarray]
-    RealScalarLike = Union[bool, int, float, Array, np.ndarray]
+    BoolScalarLike = bool | Array | np.ndarray
+    FloatScalarLike = float | Array | np.ndarray
+    IntScalarLike = int | Array | np.ndarray
+    RealScalarLike = bool | int | float | Array | np.ndarray
 else:
     BoolScalarLike = Bool[ArrayLike, ""]
     FloatScalarLike = Float[ArrayLike, ""]