From 47c13e92fdcafbc2b997f83a20315cbc1edf9d55 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Mon, 13 Jan 2025 10:03:35 -0800
Subject: [PATCH] Add test against JAX nightly

This will catch potential incompatibilities introduced in JAX before they make it into a release.

PiperOrigin-RevId: 715011865
---
 .github/workflows/ci.yml | 17 ++++++++++++++---
 test.sh                  |  2 ++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 63104b6..2f927eb 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -8,6 +8,14 @@ on:
   schedule:
     - cron: '30 2 * * *'
 
+permissions:
+  contents: read
+  actions: write  # to cancel previous workflows
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
 jobs:
   build-and-test:
     name: "Python ${{ matrix.python-version }} on ${{ matrix.os }} jax=${{ matrix.jax-version}}"
@@ -21,11 +29,14 @@ jobs:
         include:
           - python-version: "3.9"
             os: "ubuntu-latest"
-            jax-version: "0.4.27"  # Keep this in sync with version in pyproject.toml
+            jax-version: "0.4.27"  # Keep this in sync with version in requirements.txt
+          - python-version: "3.12"
+            os: "ubuntu-latest"
+            jax-version: "nightly"
 
     steps:
-      - uses: "actions/checkout@v2"
-      - uses: "actions/setup-python@v4"
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
         with:
             python-version: "${{ matrix.python-version }}"
             cache: "pip"
diff --git a/test.sh b/test.sh
index dd701a5..5c3f961 100755
--- a/test.sh
+++ b/test.sh
@@ -37,6 +37,8 @@ if [ "$JAX_VERSION" = "" ]; then
   : # use version installed in requirements above
 elif [ "$JAX_VERSION" = "newest" ]; then
   pip install -U jax jaxlib
+elif [ "$JAX_VERSION" = "nightly" ]; then
+  pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
 else
   pip install "jax==${JAX_VERSION}" "jaxlib==${JAX_VERSION}"
 fi