From 44efae32bdc5704af200c1e86e6c2601bf0bae1f Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Fri, 10 Jan 2025 14:03:55 -0800
Subject: [PATCH] CI: add scheduled test against nightly JAX releases

---
 .github/workflows/jax_nightly.yml | 52 +++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 .github/workflows/jax_nightly.yml

diff --git a/.github/workflows/jax_nightly.yml b/.github/workflows/jax_nightly.yml
new file mode 100644
index 0000000000..7beb35e381
--- /dev/null
+++ b/.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.11"]
+    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 }}
\ No newline at end of file