feat: benchmarks

.github/workflows/ci.yml

@@ -125,3 +125,36 @@ jobs:
         uses: codecov/[email protected]
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
+
+  benchmark:
+    name: Benchmark
+    runs-on: ubuntu-latest
+    needs: [smoke]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          allow-prereleases: true
+      # - name: Install dependencies
+      #   run: python -m pip install '.[test]'
+      # - name: Run benchmarks
+      #   uses: CodSpeedHQ/action@v3
+      #   with:
+      #     token: ${{ secrets.CODSPEED_TOKEN }}
+      #     run: python -m pytest tests/benchmark --codspeed
+
+      - name: Install Nox
+        run: pip install nox
+
+      - name: Install dependencies
+        run: nox --sessions benchmark --install-only
+
+      - name: Run the action
+        uses: CodSpeedHQ/action@v3
+        with:
+          run:
+            nox --sessions benchmark --reuse-existing-virtualenvs --no-install
+          token: ${{ secrets.CODSPEED_TOKEN }}

noxfile.py

@@ -39,7 +39,7 @@ def pylint(session: nox.Session) -> None:
 @nox.session
 def tests(session: nox.Session) -> None:
     """Run the unit and regular tests."""
-    session.install("-e", ".[test]")
+    session.install(".[test]")
     os.environ["GALAX_ENABLE_RUNTIME_TYPECHECKS"] = "1"  # TODO: set in a better way
     session.run("pytest", *session.posargs)
 
@@ -132,3 +132,10 @@ def build(session: nox.Session) -> None:
 
     session.install("build")
     session.run("python", "-m", "build")
+
+
+@nox.session
+def benchmark(session: nox.Session) -> None:
+    """Run the benchmark tests."""
+    session.install(".[test]")
+    session.run("python", "-m", "pytest", "tests/benchmark", "--codspeed")

pyproject.toml

@@ -45,7 +45,7 @@
     "unxt>=0.20.2",
     "xmmutablemap>=0.1",
     "zeroth>=0.1",
-]
+  ]
 
   [project.optional-dependencies]
     # All runtime dependencies
@@ -75,6 +75,7 @@
       "nox>=2024.10.9",
       "pre-commit>=4.0.1",
       "pytest-arraydiff>=0.6.1",
+      "pytest-codspeed",
       "pytest-cov>=5",
       "pytest>=8.3",
       "sybil>=8.0.0",

tests/benchmark/__init__.py

@@ -0,0 +1,5 @@
+"""Benchmark tests."""
+
+import jax
+
+jax.config.update("jax_enable_x64", True)  # noqa: FBT003