Add optional JAX backend, migrate Poetry->uv, and benchmark JAX vs NumPy

.github/actions/build/action.yml

@@ -4,43 +4,31 @@ description: Build the package
 runs:
    using: "composite"
    steps:
-    - uses: actions/checkout@v3
-    
-    - name: Set up Python
-      uses: actions/setup-python@v3
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
         python-version: '3.14'
 
-    - name: Install dependencies
-      run: |
-        python3 -m pip install --upgrade pip
-        python3 -m pip install flake8 pytest
-        python3 -m pip install poetry
-        python3 -m pip install twine
-      shell: bash  
-
-    - name: Lint with flake8
+    - name: Lint with ruff
       run: |
         # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings.
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=88 --statistics
+        uvx ruff check . --select=E9,F63,F7,F82 --output-format=full
+        # exit-zero treats all findings (line length + complexity) as warnings
+        uvx ruff check . --select=E501,C901 --line-length=88 --exit-zero --statistics
       shell: bash
 
-    - name: Install CVMatrix dependencies
-      run: poetry install
-      shell: bash
-
     - name: Build a binary wheel and a source tarball
-      run: poetry build
+      run: uv build
       shell: bash
-    
+
     - name: Check the distribution with twine
-      run: twine check dist/*
+      run: uvx twine check dist/*
       shell: bash
-      
+
     - name: Store the distribution packages
       uses: actions/upload-artifact@v4
       with:
         name: python-package-distributions
-        path: dist/
+        path: dist/

.github/actions/test/action.yml

@@ -5,31 +5,33 @@ inputs:
   PYTHON_VERSION:
     description: Test Python version
     required: true
+  INSTALL_JAX:
+    description: Whether to install the optional JAX backend extra (true/false)
+    required: false
+    default: 'true'
 
 runs:
   using: "composite"
   steps:
-    - name: Set up Python
-      uses: actions/setup-python@v5
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
       with:
         python-version: ${{ inputs.PYTHON_VERSION }}
 
-    - name: Install Poetry
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install poetry
-      shell: bash
-
-    - name: Install project dependencies
-      run: poetry install
+    - name: Install project dependencies (with JAX backend)
+      if: ${{ inputs.INSTALL_JAX == 'true' }}
+      run: uv sync --extra jax --group dev
       shell: bash
 
-    - name: Install additional dependencies for testing
-      run: |
-        poetry add --group dev pandas pytest flake8 pytest-cov typeguard
+    - name: Install project dependencies (numpy-only, without JAX)
+      if: ${{ inputs.INSTALL_JAX != 'true' }}
+      run: uv sync --group dev
       shell: bash
 
+    # The same test suite runs in both modes. With the JAX backend installed it also
+    # exercises the parametrized backend="jax" tests; without JAX those parametrizations
+    # are skipped (pytest.importorskip), verifying the numpy-only path.
     - name: Run tests
       run: |
-        poetry run pytest tests --doctest-modules --junitxml=junit/test-results.xml --cov=cvmatrix/ --cov-report=xml --typeguard-packages=cvmatrix/
+        uv run pytest tests --doctest-modules --junitxml=junit/test-results.xml --cov=cvmatrix/ --cov-report=xml --typeguard-packages=cvmatrix/
       shell: bash

.github/workflows/pull_request_package_workflow.yml

@@ -8,7 +8,7 @@ on:
             - 'tests/**'
             - '.github/**'
             - 'pyproject.toml'
-            - 'poetry.lock'
+            - 'uv.lock'
 
 jobs:
     build_package:

.github/workflows/pull_request_test_workflow.yml

@@ -8,32 +8,57 @@ on:
       - 'tests/**'
       - '.github/**'
       - 'pyproject.toml'
-      - 'poetry.lock'
+      - 'uv.lock'
 
 jobs:
-  test_package:
+  # Full test suite with the JAX backend installed (also runs backend="jax" tests)
+  test_package_with_jax:
     runs-on: ${{ matrix.os }}
-    strategy:  
-      fail-fast: false  
-      matrix:  
+    strategy:
+      fail-fast: false
+      matrix:
         os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ["3.11", "3.12", "3.13", "3.14"]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test
         with:
           PYTHON_VERSION: ${{ matrix.python-version }}
-
+          INSTALL_JAX: 'true'
+
+      - name: Upload coverage to Coveralls
+        uses: coverallsapp/github-action@v2
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          path-to-lcov: coverage.xml
+          flag-name: ${{ matrix.os }}-python-${{ matrix.python-version }}-jax
+          parallel: true
+
+  # Numpy-only installation (without JAX); backend="jax" tests are skipped
+  test_package_numpy_only:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+        python-version: ["3.11", "3.12", "3.13", "3.14"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/test
+        with:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+          INSTALL_JAX: 'false'
+
       - name: Upload coverage to Coveralls
         uses: coverallsapp/github-action@v2
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           path-to-lcov: coverage.xml
-          flag-name: ${{ matrix.os }}-python-${{ matrix.python-version }}
+          flag-name: ${{ matrix.os }}-python-${{ matrix.python-version }}-numpy-only
           parallel: true
 
   finish:
-    needs: test_package
+    needs: [test_package_with_jax, test_package_numpy_only]
     runs-on: ubuntu-latest
     steps:
       - name: Finish Coveralls

.github/workflows/test_workflow.yml

@@ -8,32 +8,57 @@ on:
             - 'tests/**'
             - '.github/**'
             - 'pyproject.toml'
-            - 'poetry.lock'
+            - 'uv.lock'
 
 jobs:
-    test_package:
-        runs-on: ${{ matrix.os }}  
-        strategy:  
-            fail-fast: false  
-            matrix:  
+    # Full test suite with the JAX backend installed (also runs backend="jax" tests)
+    test_package_with_jax:
+        runs-on: ${{ matrix.os }}
+        strategy:
+            fail-fast: false
+            matrix:
                 os: [ubuntu-latest, windows-latest, macos-latest]
                 python-version: ["3.11", "3.12", "3.13", "3.14"]
         steps:
             - uses: actions/checkout@v4
             - uses: ./.github/actions/test
               with:
                     PYTHON_VERSION: ${{ matrix.python-version }}
-
+                    INSTALL_JAX: 'true'
+
+            - name: Upload coverage to Coveralls
+              uses: coverallsapp/github-action@v2
+              with:
+                github-token: ${{ secrets.GITHUB_TOKEN }}
+                path-to-lcov: coverage.xml
+                flag-name: ${{ matrix.os }}-python-${{ matrix.python-version }}-jax
+                parallel: true
+
+    # Numpy-only installation (without JAX); backend="jax" tests are skipped
+    test_package_numpy_only:
+        runs-on: ${{ matrix.os }}
+        strategy:
+            fail-fast: false
+            matrix:
+                os: [ubuntu-latest, windows-latest, macos-latest]
+                python-version: ["3.11", "3.12", "3.13", "3.14"]
+        steps:
+            - uses: actions/checkout@v4
+            - uses: ./.github/actions/test
+              with:
+                    PYTHON_VERSION: ${{ matrix.python-version }}
+                    INSTALL_JAX: 'false'
+
             - name: Upload coverage to Coveralls
               uses: coverallsapp/github-action@v2
               with:
                 github-token: ${{ secrets.GITHUB_TOKEN }}
                 path-to-lcov: coverage.xml
-                flag-name: ${{ matrix.os }}-python-${{ matrix.python-version }}
+                flag-name: ${{ matrix.os }}-python-${{ matrix.python-version }}-numpy-only
                 parallel: true
 
     finish:
-        needs: test_package
+        needs: [test_package_with_jax, test_package_numpy_only]
         runs-on: ubuntu-latest
         steps:
             - name: Finish Coveralls

CHANGELOG.md

@@ -5,7 +5,19 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [4.0.0] - 2025
+## [3.2.0] - 2026-06-29
+
+### Added
+- Optional **JAX backend** for `CVMatrix` via `backend="jax"` (a `Literal["numpy", "jax"]`; default remains `"numpy"`). All array operations are routed through a resolved array namespace (`numpy` or `jax.numpy`), so the per-fold `training_XTX`/`training_XTY`/`training_XTX_XTY`/`training_statistics` computations can be traced by `jax.jit` and batched with `jax.vmap` on CPU/GPU/TPU. Install with `cvmatrix[jax]`.
+- Backend-neutral, trace-safe standard-deviation clamping (`maximum`/`where` instead of boolean-mask assignment). Degenerate-fold validation (`ValueError`s) still fires for eager (numpy or concrete jax) execution; under `jax.jit`/`jax.vmap` tracing the check is deferred to a host-side pre-flight by the caller.
+
+### Changed
+- Switched build and development tooling from Poetry to [uv](https://docs.astral.sh/uv/): the project now uses PEP 621 `[project]` metadata, the `hatchling` build backend, and a PEP 735 `[dependency-groups]` dev group (`poetry.lock` is replaced by `uv.lock`). The published package, its runtime dependencies, and the `jax` extra are unchanged.
+
+### Notes
+- The numpy backend is byte-identical to previous releases (verified against the existing test suite) and performance-neutral.
+
+## [3.1.6] - 2025
 ### Changed
 - `CVMatrix` and `Partitioner` can now be imported with `from cvmatrix import CVMatrix, Partitioner` or `import cvmatrix` followed by `cvmatrix.CVMatrix` and `cvmatrix.Partitioner`.
 

CONTRIBUTING.md

@@ -44,7 +44,7 @@ If you feel like you've made a valuable contribution, but you don't know how to
 
 ## Build from source
 
-CVMatrix uses [poetry](https://python-poetry.org/) to manage its dependencies and packaging. To build the package from source, follow these steps:
+CVMatrix uses [uv](https://docs.astral.sh/uv/) to manage its dependencies and packaging. To build the package from source, follow these steps:
 
 1. Clone the repository:
 
@@ -58,57 +58,50 @@ CVMatrix uses [poetry](https://python-poetry.org/) to manage its dependencies an
     cd CVMatrix
     ```
 
-3. Install poetry and twine:
+3. [Install uv](https://docs.astral.sh/uv/getting-started/installation/), e.g.:
 
     ```shell
-    pip3 install poetry
-    pip3 install twine
+    curl -LsSf https://astral.sh/uv/install.sh | sh
     ```
 
-4. Install the dependencies:
+4. Build the package. uv builds in an isolated environment, so no separate dependency-install step is needed:
 
     ```shell
-    poetry install
+    uv build
     ```
 
-5. Build the package:
+5. Check the package with twine:
 
     ```shell
-    poetry build
-    ```
-
-6. Check the package with twine:
-
-    ```shell
-    twine check dist/*
+    uvx twine check dist/*
     ```
 
 ## Run the test suite
 
 To run the test suite, follow these steps:
 
-1. Install poetry:
+1. [Install uv](https://docs.astral.sh/uv/getting-started/installation/), e.g.:
 
     ```shell
-    pip3 install poetry
+    curl -LsSf https://astral.sh/uv/install.sh | sh
     ```
 
-2. Install the project dependencies with poetry:
+2. Install the project together with its development dependencies. To also exercise the optional JAX backend (the `backend="jax"` tests), include the `jax` extra:
 
     ```shell
-    poetry install
+    uv sync --extra jax --group dev
     ```
 
-3. Install additional dependencies for testing:
+    To test the numpy-only installation instead (the `backend="jax"` tests are skipped automatically), omit the extra:
 
     ```shell
-    poetry add --group dev pandas pytest flake8 pytest-cov typeguard
+    uv sync --group dev
     ```
 
-4. Now, the tests can be run with the following command:
+3. Now, the tests can be run with the following command:
 
     ```shell
-    poetry run pytest tests --doctest-modules --junitxml=junit/test-results.xml --cov=cvmatrix/ --cov-report=xml --cov-report=html --typeguard-packages=cvmatrix/
+    uv run pytest tests --doctest-modules --junitxml=junit/test-results.xml --cov=cvmatrix/ --cov-report=xml --cov-report=html --typeguard-packages=cvmatrix/
     ```
 
 ## Build the documentation