Update infra (#182)

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 }}/tests ./tests
                 cp ${{ github.workspace }}/pyproject.toml ./pyproject.toml
-                python -m pip install -r ./tests/requirements.txt
-                python -m tests
+                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

@@ -1,17 +1,3 @@
-# Copyright 2023 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 name: Run tests
 
 on:
@@ -21,7 +7,8 @@ jobs:
   run-test:
     strategy:
       matrix:
-        python-version: [ "3.10", "3.12" ]
+        # must match the `language_version` in `.pre-commit-config.yaml`
+        python-version: [ 3.11 ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
@@ -37,12 +24,16 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -r ./tests/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 tests
+          pytest
+
+      - name: Check that documentation can be built.
+        run: |
+          mkdocs build

.gitignore

@@ -9,3 +9,4 @@ examples/data
 .pymon
 .idea
 .venv
+uv.lock

.pre-commit-config.yaml

@@ -22,15 +22,15 @@ repos:
         files: ^pyproject\.toml$
         additional_dependencies: ["toml-sort==0.23.1"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.7
+    rev: v0.13.0
     hooks:
       - id: ruff-format
         types_or: [ python, pyi, jupyter ]
       - id: ruff
         types_or: [ python, pyi, jupyter ]
         args: [ --fix ]
   - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.330
+    rev: v1.1.406
     hooks:
     - id: pyright
       additional_dependencies: ["jax", "equinox", "pytest"]

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/lineax.git
 cd lineax
-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 Black to format the code, and ruff to lint 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 tests/requirements.txt
-python -m tests
+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:
@@ -53,26 +45,8 @@ Finally, open a pull request on GitHub!
 Make your changes. You can then build the documentation by doing
 
 ```bash
-pip install -r docs/requirements.txt
+pip install -e '.[docs]'
 mkdocs serve
 ```
-Then doing `Control-C`, and running:
-```
-mkdocs serve
-```
-(So you run `mkdocs serve` twice.)
 
 You can then see your local copy of the documentation by navigating to `localhost:8000` in a web browser.
-
-## Contributor License Agreement
-
-Contributions to this project must be accompanied by a Contributor License
-Agreement (CLA). You (or your employer) retain the copyright to your
-contribution; this simply gives us permission to use and redistribute your
-contributions as part of the project. Head over to
-<https://cla.developers.google.com/> to see your current agreements on file or
-to sign a new one.
-
-You generally only need to submit a CLA once, so if you've already submitted one
-(even if it was for a different project), you probably don't need to do it
-again.

lineax/_solver/gmres.py

@@ -410,7 +410,7 @@ def _normalise(
             eps = jnp.finfo(norm.dtype).eps
         else:
             eps = jnp.astype(eps, norm.dtype)
-        breakdown = norm < eps
+        breakdown = norm < eps  # pyright: ignore
         safe_norm = jnp.where(breakdown, jnp.inf, norm)
         with jax.numpy_dtype_promotion("standard"):
             x_normalised = (x**ω / safe_norm).ω

pyproject.toml

@@ -31,15 +31,27 @@ urls = {repository = "https://github.com/google/lineax"}
 version = "0.0.8"
 
 [project.optional-dependencies]
+dev = [
+  "pre-commit"
+]
 docs = [
-  "hippogriffe==0.2.0",
+  "hippogriffe==0.2.2",
+  "griffe==1.7.3",
   "mkdocs==1.6.1",
   "mkdocs-include-exclude-files==0.1.0",
-  "mkdocs-ipynb==0.1.0",
+  "mkdocs-ipynb==0.1.1",
   "mkdocs-material==9.6.7",
-  "mkdocstrings[python]==0.28.3",
+  "mkdocstrings==0.28.3",
+  "mkdocstrings-python==1.16.8",
   "pymdown-extensions==10.14.3"
 ]
+tests = [
+  "beartype",
+  "equinox",
+  "pytest",
+  "pytest-xdist",
+  "jaxlib"
+]
 
 [tool.hatch.build]
 include = ["lineax/*"]
@@ -53,14 +65,20 @@ addopts = "--jaxtyping-packages=lineax,beartype.beartype(conf=beartype.BeartypeC
 
 [tool.ruff]
 extend-include = ["*.ipynb"]
-fixable = ["I001", "F401", "UP"]
-ignore = ["E402", "E721", "E731", "E741", "F722", "UP038"]
-ignore-init-module-imports = true
-select = ["E", "F", "I001", "UP"]
 src = []
 
 [tool.ruff.isort]
 combine-as-imports = true
 extra-standard-library = ["typing_extensions"]
 lines-after-imports = 2
 order-by-type = false
+
+[tool.ruff.lint]
+fixable = ["I001", "F401", "UP"]
+ignore = ["E402", "E721", "E731", "E741", "F722", "UP038"]
+select = ["E", "F", "I001", "UP"]
+
+[tool.ruff.lint.flake8-import-conventions.extend-aliases]
+"collections" = "co"
+"functools" = "ft"
+"itertools" = "it"