CI: add scheduled test against nightly JAX releases

.github/workflows/flax_test.yml

@@ -30,7 +30,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
       - uses: pre-commit/[email protected]
   commit-count:
     name: Check commit count

.github/workflows/jax_nightly.yml

@@ -0,0 +1,52 @@
+name: CI - with JAX nightly
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+on:
+  schedule:
+    - cron: "0 12 * * *" # Daily at 12:00 UTC
+  workflow_dispatch: # allows triggering the workflow run manually
+  pull_request:  # Automatically trigger on pull requests affecting this file
+    branches:
+      - main
+    paths:
+      - '**workflows/jax_nightly.yml'
+
+jobs:
+  jax-nightly:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      issues: write  # for failed-build-issue
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12"]
+    steps:
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+    - name: Set up Python ${{ matrix.python-version }}
+      id: setup_python
+      uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Setup uv
+      uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a  # v5.1.0
+      with:
+        version: "0.3.0"
+    - name: Install dependencies
+      run: |
+        uv sync --extra all --extra testing --extra docs
+    - name: Install JAX
+      run: |
+        uv pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+    - name: Run test suite
+      if: success()
+      run: |
+        uv run tests/run_all_tests.sh --only-pytest
+    - name: Notify failed build
+      uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
+      if: failure() && github.event.pull_request == null
+      with:
+        github-token: ${{ secrets.GITHUB_TOKEN }}