diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index d8651529..205f4e05 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -8,3 +8,7 @@ updates:
     directory: "/"
     schedule:
       interval: "weekly"
+  - package-ecosystem: "uv"
+    directory: "/"
+    schedule:
+      interval: "weekly"
diff --git a/.github/workflows/uv_lock_autoupdate.yml b/.github/workflows/uv_lock_autoupdate.yml
deleted file mode 100644
index 55b6f588..00000000
--- a/.github/workflows/uv_lock_autoupdate.yml
+++ /dev/null
@@ -1,139 +0,0 @@
-name: Autoupdate uv lock file
-
-on:
-  schedule:
-    # Every Sunday at 00:00.
-    - cron: "0 0 * * 0"
-
-jobs:
-  uv_lock_autoupdate:
-    env:
-      lockfile_path: uv.lock
-      requirements_doc_path: requirements-doc.txt
-      update_message_path: /tmp/message.txt
-      update_branch: chore/uv-lock-autoupdate
-    runs-on: ubuntu-latest
-    steps:
-      - name: Generate a token
-        id: generate-token
-        uses: actions/create-github-app-token@v2
-        with:
-          app-id: ${{ vars.WRITE_CONTENTS_PR_APP }}
-          private-key: ${{ secrets.WRITE_CONTENTS_PR_KEY }}
-      - uses: actions/checkout@v5
-        with:
-          ref: main
-      - name: Create or checkout update branch
-        id: create_branch
-        env:
-          GH_TOKEN: ${{ steps.generate-token.outputs.token }}
-        run: |
-          export pr_number=$( gh pr view ${{ env.update_branch }} --json number --jq '.number' )
-          export pr_state=$( gh pr view ${{ env.update_branch }} --json state --jq '.state' )
-          echo "pr_number=$pr_number" >> "$GITHUB_OUTPUT"
-          echo "pr_state=$pr_state" >> "$GITHUB_OUTPUT"
-
-          if git fetch origin ${{ env.update_branch }}; then
-            # If branch wasn't deleted after merge, do so here
-            if [ "$pr_state" = "MERGED" ]; then
-              git push -d origin ${{ env.update_branch }}
-              git checkout -b ${{ env.update_branch }}
-              git push origin ${{ env.update_branch }}
-            fi
-            git checkout ${{ env.update_branch }}
-          else
-            git checkout -b ${{ env.update_branch }}
-            git push origin ${{ env.update_branch }}
-          fi
-
-      # uv setup and lockfile upgrading steps
-      - name: Set up uv
-        uses: astral-sh/setup-uv@v6
-        with:
-          version: "latest"
-          python-version: 3.13
-      - name: Store hash of baseline files for comparison
-        id: old_hash
-        run: |
-          echo "lockfile=$( sha256sum $lockfile_path )" >> $GITHUB_OUTPUT
-          echo "requirements_doc=$( sha256sum $requirements_doc_path )" >> $GITHUB_OUTPUT
-      - name: Overwrite config on branch with version from main
-        run: |
-          git checkout main ${{ env.lockfile_path }}
-          git checkout main ${{ env.requirements_doc_path }}
-      - name: Store hash of main files for comparison
-        id: main_hash
-        run: |
-          echo "lockfile=$( sha256sum $lockfile_path )" >> $GITHUB_OUTPUT
-          echo "requirements_doc=$( sha256sum $requirements_doc_path )" >> $GITHUB_OUTPUT
-      - name: Run `uv lock --upgrade` on main lockfile
-        id: autoupdate
-        run: |
-          uv lock --upgrade > ${{ env.update_message_path }}
-          cat ${{ env.update_message_path }}
-          sed -i "/updated/!d" ${{ env.update_message_path }}
-      - name: Update requirements-docs.txt
-        run: >
-          uv export
-          --frozen
-          --format requirements-txt
-          --no-hashes
-          --no-dev
-          --no-editable
-          --extra doc
-          --output-file requirements-doc.txt
-      - name: Store hash of new files for comparison
-        id: new_hash
-        run: |
-          echo "lockfile=$( sha256sum $lockfile_path )" >> $GITHUB_OUTPUT
-          echo "requirements_doc=$( sha256sum $requirements_doc_path )" >> $GITHUB_OUTPUT
-      # Commit authoring and pull-request creation/updating
-      - name: Commit (with signature) lockfile
-        id: commit_lockfile
-        if: steps.old_hash.outputs.lockfile != steps.new_hash.outputs.lockfile
-        env:
-          GH_TOKEN: ${{ steps.generate-token.outputs.token }}
-        run: |
-          export message="chore(deps): autoupdate uv.lock"
-          export sha=$( git rev-parse ${{ env.update_branch }}:${{ env.lockfile_path }} )
-          base64 -i ${{ env.lockfile_path }} > $RUNNER_TEMP/base64_lockfile
-          gh api --method PUT /repos/:owner/:repo/contents/${{ env.lockfile_path }} \
-            --field message="$message" \
-            --field content="@$RUNNER_TEMP/base64_lockfile" \
-            --field branch=${{ env.update_branch }} \
-            --field sha="$sha"
-      - name: Commit (with signature) new requirements-doc
-        id: commit_requirements_doc
-        if: steps.old_hash.outputs.requirements_doc != steps.new_hash.outputs.requirements_doc
-        env:
-          GH_TOKEN: ${{ steps.generate-token.outputs.token }}
-        run: |
-          export message="chore(deps): autoupdate requirements-doc"
-          export sha=$( git rev-parse ${{ env.update_branch }}:${{ env.requirements_doc_path }} )
-          base64 -i ${{ env.requirements_doc_path }} > $RUNNER_TEMP/base64_requirements_doc
-          gh api --method PUT /repos/:owner/:repo/contents/${{ env.requirements_doc_path }} \
-            --field message="$message" \
-            --field content="@$RUNNER_TEMP/base64_requirements_doc" \
-            --field branch=${{ env.update_branch }} \
-            --field sha="$sha"
-      - name: Create or update lockfile-autoupdate pull-request
-        # note that we don't check whether requirements-doc was updated, as
-        # requirements-doc should only change if the lockfile also changes
-        if: steps.commit_lockfile.conclusion == 'success' || ( steps.main_lockfile.outputs.hash != steps.new_lockfile.outputs.hash )
-        env:
-          GH_TOKEN: ${{ steps.generate-token.outputs.token }}
-        run: |
-          export title="chore(deps): autoupdate uv.lock"
-          export body=$( cat ${{ env.update_message_path }} )
-          export pr_number=${{ steps.create_branch.outputs.pr_number }}
-          export pr_state=${{ steps.create_branch.outputs.pr_state }}
-
-          # If the PR is closed, can it be reopened, or is the PR already open?
-          if ( [ "$pr_state" = "CLOSED" ] && gh pr reopen $pr_number ) || [ "$pr_state" = "OPEN" ]; then
-            gh api --method PATCH /repos/:owner/:repo/pulls/$pr_number \
-              --field title="$title" \
-              --field body="$body"
-          else
-            # If a PR doesn't already exist, and no previous PR can be reopened, create a new PR.
-            gh pr create -t "$title" -b "$body" -l dependencies -B main -H ${{ env.update_branch }}
-          fi
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b41017e1..5a82244e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -46,18 +46,6 @@ repos:
     hooks:
       # Keep lock file up to date
       - id: uv-lock
-      # Keep requirements-doc up to date
-      - id: uv-export
-        args: [
-          "--frozen",
-          "--quiet",
-          "--format", "requirements-txt",
-          "--no-hashes",
-          "--no-dev",
-          "--no-editable",
-          "--extra", "doc",
-          "--output-file", "requirements-doc.txt",
-        ]
   - repo: local
     hooks:
       # Check files are valid UTF-8
diff --git a/requirements-doc.txt b/requirements-doc.txt
deleted file mode 100644
index 078ad963..00000000
--- a/requirements-doc.txt
+++ /dev/null
@@ -1,447 +0,0 @@
-# This file was autogenerated by uv via the following command:
-#    uv export --frozen --format requirements-txt --no-hashes --no-dev --no-editable --extra doc --output-file requirements-doc.txt
-.
-absl-py==2.3.1
-    # via
-    #   chex
-    #   optax
-    #   orbax-checkpoint
-accessible-pygments==0.0.5
-    # via furo
-aiofiles==24.1.0 ; python_full_version >= '3.10'
-    # via orbax-checkpoint
-alabaster==0.7.16 ; python_full_version < '3.10'
-    # via sphinx
-alabaster==1.0.0 ; python_full_version >= '3.10'
-    # via sphinx
-apeye==1.4.1
-    # via sphinx-toolbox
-apeye-core==1.1.5
-    # via apeye
-attrs==25.3.0
-    # via
-    #   jsonschema
-    #   referencing
-    #   sphobjinv
-autodocsumm==0.2.14
-    # via sphinx-toolbox
-babel==2.17.0
-    # via sphinx
-beautifulsoup4==4.13.4
-    # via
-    #   furo
-    #   sphinx-toolbox
-cachecontrol==0.14.3
-    # via sphinx-toolbox
-certifi==2025.8.3
-    # via
-    #   requests
-    #   sphinx-prompt
-    #   sphobjinv
-charset-normalizer==3.4.3
-    # via requests
-chex==0.1.90
-    # via optax
-colorama==0.4.6 ; sys_platform == 'win32'
-    # via
-    #   sphinx
-    #   tqdm
-cssutils==2.11.1
-    # via dict2css
-dict2css==0.3.0.post1
-    # via sphinx-toolbox
-docutils==0.21.2
-    # via
-    #   pybtex-docutils
-    #   sphinx
-    #   sphinx-prompt
-    #   sphinx-tabs
-    #   sphinx-toolbox
-    #   sphinxcontrib-bibtex
-domdf-python-tools==3.10.0
-    # via
-    #   apeye
-    #   apeye-core
-    #   dict2css
-    #   sphinx-toolbox
-equinox==0.11.10 ; python_full_version < '3.10'
-    # via coreax
-equinox==0.13.0 ; python_full_version >= '3.10'
-    # via coreax
-etils==1.5.2 ; python_full_version < '3.10'
-    # via
-    #   optax
-    #   orbax-checkpoint
-etils==1.13.0 ; python_full_version >= '3.10'
-    # via orbax-checkpoint
-filelock==3.18.0
-    # via
-    #   cachecontrol
-    #   sphinx-toolbox
-flax==0.8.5 ; python_full_version < '3.10'
-    # via coreax
-flax==0.10.7 ; python_full_version == '3.10.*'
-    # via coreax
-flax==0.11.1 ; python_full_version >= '3.11'
-    # via coreax
-fsspec==2025.7.0
-    # via etils
-furo==2025.7.19
-    # via coreax
-html5lib==1.1
-    # via sphinx-toolbox
-humanize==4.12.3
-    # via orbax-checkpoint
-idna==3.10
-    # via
-    #   apeye-core
-    #   requests
-    #   sphinx-prompt
-imagesize==1.4.1
-    # via sphinx
-importlib-metadata==8.7.0 ; python_full_version < '3.10'
-    # via
-    #   jax
-    #   pybtex
-    #   sphinx
-    #   sphinxcontrib-bibtex
-importlib-resources==6.5.2
-    # via etils
-jax==0.4.30 ; python_full_version < '3.10'
-    # via
-    #   chex
-    #   coreax
-    #   equinox
-    #   flax
-    #   jaxopt
-    #   optax
-    #   orbax-checkpoint
-jax==0.6.2 ; python_full_version == '3.10.*'
-    # via
-    #   chex
-    #   coreax
-    #   equinox
-    #   flax
-    #   jaxopt
-    #   optax
-    #   orbax-checkpoint
-jax==0.7.0 ; python_full_version >= '3.11'
-    # via
-    #   chex
-    #   coreax
-    #   equinox
-    #   flax
-    #   jaxopt
-    #   optax
-    #   orbax-checkpoint
-jaxlib==0.4.30 ; python_full_version < '3.10'
-    # via
-    #   chex
-    #   jax
-    #   jaxopt
-    #   optax
-    #   orbax-checkpoint
-jaxlib==0.6.2 ; python_full_version == '3.10.*'
-    # via
-    #   chex
-    #   jax
-    #   jaxopt
-    #   optax
-jaxlib==0.7.0 ; python_full_version >= '3.11'
-    # via
-    #   chex
-    #   jax
-    #   jaxopt
-    #   optax
-jaxopt==0.8.5
-    # via coreax
-jaxtyping==0.2.36 ; python_full_version < '3.10'
-    # via
-    #   coreax
-    #   equinox
-jaxtyping==0.3.2 ; python_full_version >= '3.10'
-    # via
-    #   coreax
-    #   equinox
-jinja2==3.1.6
-    # via
-    #   sphinx
-    #   sphinx-jinja2-compat
-    #   sphinx-prompt
-joblib==1.5.1
-    # via scikit-learn
-jsonschema==4.25.0
-    # via sphobjinv
-jsonschema-specifications==2025.4.1
-    # via jsonschema
-latexcodec==3.0.1
-    # via pybtex
-markdown-it-py==3.0.0
-    # via rich
-markupsafe==3.0.2
-    # via
-    #   jinja2
-    #   sphinx-jinja2-compat
-mdurl==0.1.2
-    # via markdown-it-py
-ml-dtypes==0.5.3
-    # via
-    #   jax
-    #   jaxlib
-    #   tensorstore
-more-itertools==10.7.0
-    # via cssutils
-msgpack==1.1.1
-    # via
-    #   cachecontrol
-    #   flax
-    #   orbax-checkpoint
-natsort==8.4.0
-    # via domdf-python-tools
-nest-asyncio==1.6.0
-    # via orbax-checkpoint
-numpy==2.0.2 ; python_full_version < '3.10'
-    # via
-    #   chex
-    #   flax
-    #   jax
-    #   jaxlib
-    #   jaxopt
-    #   ml-dtypes
-    #   optax
-    #   orbax-checkpoint
-    #   scikit-learn
-    #   scipy
-    #   tensorstore
-numpy==2.2.6 ; python_full_version >= '3.10'
-    # via
-    #   chex
-    #   flax
-    #   jax
-    #   jaxlib
-    #   jaxopt
-    #   ml-dtypes
-    #   optax
-    #   orbax-checkpoint
-    #   scikit-learn
-    #   scipy
-    #   tensorstore
-    #   treescope
-opt-einsum==3.4.0
-    # via jax
-optax==0.2.4 ; python_full_version < '3.10'
-    # via
-    #   coreax
-    #   flax
-optax==0.2.5 ; python_full_version >= '3.10'
-    # via
-    #   coreax
-    #   flax
-orbax-checkpoint==0.6.4 ; python_full_version < '3.10'
-    # via flax
-orbax-checkpoint==0.11.21 ; python_full_version >= '3.10'
-    # via flax
-packaging==25.0
-    # via sphinx
-platformdirs==4.3.8
-    # via apeye
-protobuf==6.31.1
-    # via orbax-checkpoint
-pybtex==0.25.1
-    # via
-    #   pybtex-docutils
-    #   sphinxcontrib-bibtex
-pybtex-docutils==1.0.3
-    # via sphinxcontrib-bibtex
-pygments==2.19.2
-    # via
-    #   accessible-pygments
-    #   furo
-    #   rich
-    #   sphinx
-    #   sphinx-prompt
-    #   sphinx-tabs
-pyyaml==6.0.2
-    # via
-    #   flax
-    #   orbax-checkpoint
-    #   pybtex
-referencing==0.36.2
-    # via
-    #   jsonschema
-    #   jsonschema-specifications
-requests==2.32.4
-    # via
-    #   apeye
-    #   cachecontrol
-    #   sphinx
-    #   sphinx-prompt
-rich==14.1.0
-    # via flax
-roman-numerals-py==3.1.0 ; python_full_version >= '3.11'
-    # via sphinx
-rpds-py==0.27.0
-    # via
-    #   jsonschema
-    #   referencing
-ruamel-yaml==0.18.14
-    # via sphinx-toolbox
-ruamel-yaml-clib==0.2.12 ; python_full_version < '3.14' and platform_python_implementation == 'CPython'
-    # via ruamel-yaml
-scikit-learn==1.6.1 ; python_full_version < '3.10'
-    # via coreax
-scikit-learn==1.7.1 ; python_full_version >= '3.10'
-    # via coreax
-scipy==1.13.1 ; python_full_version < '3.10'
-    # via
-    #   jax
-    #   jaxlib
-    #   jaxopt
-    #   scikit-learn
-scipy==1.15.3 ; python_full_version == '3.10.*'
-    # via
-    #   jax
-    #   jaxlib
-    #   jaxopt
-    #   scikit-learn
-scipy==1.16.1 ; python_full_version >= '3.11'
-    # via
-    #   jax
-    #   jaxlib
-    #   jaxopt
-    #   scikit-learn
-setuptools==80.9.0 ; python_full_version >= '3.12'
-    # via chex
-simplejson==3.20.1 ; python_full_version >= '3.10'
-    # via orbax-checkpoint
-six==1.17.0
-    # via html5lib
-snowballstemmer==3.0.1
-    # via sphinx
-soupsieve==2.7
-    # via beautifulsoup4
-sphinx==7.4.7 ; python_full_version < '3.10'
-    # via
-    #   autodocsumm
-    #   coreax
-    #   furo
-    #   sphinx-autodoc-typehints
-    #   sphinx-basic-ng
-    #   sphinx-prompt
-    #   sphinx-tabs
-    #   sphinx-toolbox
-    #   sphinxcontrib-bibtex
-sphinx==8.1.3 ; python_full_version == '3.10.*'
-    # via
-    #   autodocsumm
-    #   coreax
-    #   furo
-    #   sphinx-autodoc-typehints
-    #   sphinx-basic-ng
-    #   sphinx-prompt
-    #   sphinx-tabs
-    #   sphinx-toolbox
-    #   sphinxcontrib-bibtex
-sphinx==8.2.3 ; python_full_version >= '3.11'
-    # via
-    #   autodocsumm
-    #   coreax
-    #   furo
-    #   sphinx-autodoc-typehints
-    #   sphinx-basic-ng
-    #   sphinx-prompt
-    #   sphinx-tabs
-    #   sphinx-toolbox
-    #   sphinxcontrib-bibtex
-sphinx-autodoc-typehints==2.3.0 ; python_full_version < '3.10'
-    # via
-    #   coreax
-    #   sphinx-toolbox
-sphinx-autodoc-typehints==3.0.1 ; python_full_version == '3.10.*'
-    # via
-    #   coreax
-    #   sphinx-toolbox
-sphinx-autodoc-typehints==3.2.0 ; python_full_version >= '3.11'
-    # via
-    #   coreax
-    #   sphinx-toolbox
-sphinx-basic-ng==1.0.0b2
-    # via furo
-sphinx-jinja2-compat==0.4.1
-    # via sphinx-toolbox
-sphinx-prompt==1.8.0 ; python_full_version < '3.10'
-    # via sphinx-toolbox
-sphinx-prompt==1.9.0 ; python_full_version == '3.10.*'
-    # via sphinx-toolbox
-sphinx-prompt==1.10.0 ; python_full_version >= '3.11'
-    # via sphinx-toolbox
-sphinx-tabs==3.4.5
-    # via sphinx-toolbox
-sphinx-toolbox==4.0.0
-    # via coreax
-sphinxcontrib-applehelp==2.0.0
-    # via sphinx
-sphinxcontrib-bibtex==2.6.5
-    # via coreax
-sphinxcontrib-devhelp==2.0.0
-    # via sphinx
-sphinxcontrib-htmlhelp==2.1.0
-    # via sphinx
-sphinxcontrib-jsmath==1.0.1
-    # via sphinx
-sphinxcontrib-qthelp==2.0.0
-    # via sphinx
-sphinxcontrib-serializinghtml==2.0.0
-    # via sphinx
-sphobjinv==2.3.1.3
-    # via coreax
-standard-imghdr==3.10.14 ; python_full_version >= '3.13'
-    # via sphinx-jinja2-compat
-tabulate==0.9.0
-    # via sphinx-toolbox
-tensorstore==0.1.69 ; python_full_version < '3.10'
-    # via
-    #   flax
-    #   orbax-checkpoint
-tensorstore==0.1.76 ; python_full_version >= '3.10'
-    # via
-    #   flax
-    #   orbax-checkpoint
-threadpoolctl==3.6.0
-    # via scikit-learn
-tomli==2.2.1 ; python_full_version < '3.11'
-    # via sphinx
-toolz==1.0.0
-    # via chex
-tqdm==4.67.1
-    # via coreax
-treescope==0.1.10 ; python_full_version >= '3.10'
-    # via flax
-typing-extensions==4.14.1
-    # via
-    #   beautifulsoup4
-    #   chex
-    #   coreax
-    #   domdf-python-tools
-    #   equinox
-    #   etils
-    #   flax
-    #   jaxtyping
-    #   orbax-checkpoint
-    #   referencing
-    #   sphinx-toolbox
-urllib3==2.5.0
-    # via
-    #   requests
-    #   sphinx-prompt
-wadler-lindig==0.1.7 ; python_full_version >= '3.10'
-    # via
-    #   equinox
-    #   jaxtyping
-webencodings==0.5.1
-    # via html5lib
-zipp==3.23.0
-    # via
-    #   etils
-    #   importlib-metadata
-    #   importlib-resources