Standardised infra

Author: patrick-kidger

.github/workflows/release.yml

@@ -10,14 +10,14 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Release
-        uses: patrick-kidger/action_update_python_project@v6
+        uses: patrick-kidger/action_update_python_project@v8
         with:
             python-version: "3.11"
             test-script: |
                 cp -r ${{ github.workspace }}/test ./test
                 cp ${{ github.workspace }}/pyproject.toml ./pyproject.toml
-                python -m pip install -r ./test/requirements.txt
-                python -m test
+                uv sync --extra tests --no-install-project --inexact
+                uv run --no-sync pytest
             pypi-token: ${{ secrets.pypi_token }}
             github-user: patrick-kidger
             github-token: ${{ github.token }}

.github/workflows/run_tests.yml

@@ -23,13 +23,16 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -r ./test/requirements.txt
-
+          python -m pip install '.[dev,docs,tests]'
 
       - name: Checks with pre-commit
-        uses: pre-commit/action@v3.0.1
+        run: |
+          pre-commit run --all-files
 
       - name: Test with pytest
         run: |
-          python -m pip install .
           python -m test
+
+      - name: Check that documentation can be built.
+        run: |
+          mkdocs build

.gitignore

@@ -8,3 +8,4 @@ site/
 .pymon
 .idea/
 .venv/
+uv.lock

.pre-commit-config.yaml

@@ -8,15 +8,15 @@ repos:
         files: ^pyproject\.toml$
         additional_dependencies: ["toml-sort==0.23.1"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.2.2
+    rev: v0.13.0
     hooks:
       - id: ruff-format  # formatter
-        types_or: [ python, pyi, jupyter ]
+        types_or: [ python, pyi, jupyter, toml ]
       - id: ruff  # linter
-        types_or: [ python, pyi, jupyter ]
+        types_or: [ python, pyi, jupyter, toml ]
         args: [ --fix ]
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.350
+    rev: v1.1.405
     hooks:
     - id: pyright
       additional_dependencies: [equinox, jax, jaxtyping, optax, optimistix, lineax, pytest, typeguard==2.13.3, typing_extensions, wadler_lindig]

CONTRIBUTING.md

@@ -8,23 +8,15 @@ Contributions (pull requests) are very welcome! Here's how to get started.
 
 First fork the library on GitHub.
 
-Then clone and install the library in development mode:
+Then clone and install the library:
 
 ```bash
 git clone https://github.com/your-username-here/diffrax.git
 cd diffrax
-pip install -e .
+pip install -e '.[dev]'
+pre-commit install  # `pre-commit` is installed by `pip` on the previous line
 ```
 
-Then install the pre-commit hook:
-
-```bash
-pip install pre-commit
-pre-commit install
-```
-
-These hooks use ruff to lint and format the code, and pyright to type-check it.
-
 ---
 
 **If you're making changes to the code:**
@@ -34,8 +26,8 @@ Now make your changes. Make sure to include additional tests if necessary.
 Next verify the tests all pass:
 
 ```bash
-pip install -r test/requirements.txt
-pytest
+pip install -e '.[tests]'
+pytest  # `pytest` is installed by `pip` on the previous line.
 ```
 
 Then push your changes back to your fork of the repository:

diffrax/_adjoint.py

@@ -362,8 +362,7 @@ def loop(
         if is_unsafe_sde(terms):
             kind = "lax"
             msg = (
-                "Cannot reverse-mode autodifferentiate when using "
-                "`UnsafeBrownianPath`."
+                "Cannot reverse-mode autodifferentiate when using `UnsafeBrownianPath`."
             )
         elif max_steps is None:
             kind = "lax"

diffrax/_brownian/path.py

@@ -62,9 +62,9 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: type[
-        BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
-    ] = eqx.field(static=True)
+    levy_area: type[BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea] = (
+        eqx.field(static=True)
+    )
     key: PRNGKeyArray
 
     def __init__(

diffrax/_brownian/tree.py

@@ -235,9 +235,9 @@ class VirtualBrownianTree(AbstractBrownianPath):
     t1: RealScalarLike
     tol: RealScalarLike
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: type[
-        BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea
-    ] = eqx.field(static=True)
+    levy_area: type[BrownianIncrement | SpaceTimeLevyArea | SpaceTimeTimeLevyArea] = (
+        eqx.field(static=True)
+    )
     key: PyTree[PRNGKeyArray]
     _spline: _Spline = eqx.field(static=True)
 

diffrax/_integrate.py

@@ -3,6 +3,7 @@
 from collections.abc import Callable
 from typing import (  # noqa: UP035
     Any,
+    cast,
     get_args,
     get_origin,
     Tuple,
@@ -1164,7 +1165,10 @@ def _wrap(term):
         def _get_tols(x):
             outs = []
             for attr in ("rtol", "atol", "norm"):
-                if getattr(solver.root_finder, attr) is use_stepsize_tol:
+                if (
+                    getattr(cast(AbstractImplicitSolver, solver).root_finder, attr)
+                    is use_stepsize_tol
+                ):
                     outs.append(getattr(x, attr))
             return tuple(outs)
 

diffrax/_solution.py

@@ -10,7 +10,7 @@
 from ._path import AbstractPath
 
 
-class RESULTS(optx.RESULTS):  # pyright: ignore
+class RESULTS(optx.RESULTS):  # pyright: ignore[reportGeneralTypeIssues]
     successful = ""
     max_steps_reached = (
         "The maximum number of solver steps was reached. Try increasing `max_steps`."
@@ -121,8 +121,8 @@ class Solution(AbstractPath):
     # the structure of `subs`.
     # SaveAt(fn=...) means that `ys` will then follow with arbitrary sub-dependent
     # PyTree structures.
-    ts: PyTree[Real[Array, " ?times"], " S"] | None
-    ys: PyTree[Shaped[Array, "?times ?*shape"], "S ..."] | None
+    ts: PyTree[Real[Array, " ?times"], " S"] | None  # pyright: ignore[reportUndefinedVariable]
+    ys: PyTree[Shaped[Array, "?times ?*shape"], "S ..."] | None  # pyright: ignore
     interpolation: DenseInterpolation | None
     stats: dict[str, Any]
     result: RESULTS
@@ -133,7 +133,7 @@ class Solution(AbstractPath):
 
     def evaluate(
         self, t0: RealScalarLike, t1: RealScalarLike | None = None, left: bool = True
-    ) -> PyTree[Shaped[Array, "?*shape"], " Y"]:
+    ) -> PyTree[Shaped[Array, "?*shape"], " Y"]:  # pyright: ignore[reportUndefinedVariable]
         """If dense output was saved, then evaluate the solution at any point in the
         region of integration `self.t0` to `self.t1`.
 
@@ -153,7 +153,7 @@ def evaluate(
 
     def derivative(
         self, t: RealScalarLike, left: bool = True
-    ) -> PyTree[Shaped[Array, "?*shape"], " Y"]:
+    ) -> PyTree[Shaped[Array, "?*shape"], " Y"]:  # pyright: ignore[reportUndefinedVariable]
         r"""If dense output was saved, then calculate an **approximation** to the
         derivative of the solution at any point in the region of integration `self.t0`
         to `self.t1`.