From 11f89eb2b7600de6b4c4860549fd8f38a244fd2d Mon Sep 17 00:00:00 2001
From: Yu Li <wstcliyu@google.com>
Date: Tue, 4 Aug 2026 21:22:19 +0000
Subject: [PATCH] Add JAX 0.11.0 to test matrix. Pin github action to commit
 SHAs.

---
 .github/workflows/test.yml | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 755ce6d..e5060bf 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,6 +17,9 @@ name: Unittests
 # Allow to trigger the workflow manually (e.g. when deps changes)
 on: [push, workflow_dispatch]
 
+permissions:
+  contents: read
+
 jobs:
   unittest-job:
     name: Unittests (Python ${{ matrix.python-version }}, JAX ${{ matrix.jax-version }})
@@ -24,7 +27,10 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ['3.11', '3.12', '3.13', '3.14']
-        jax-version: ['0.8.3', '0.9.0', '0.9.1', '0.9.2', '0.10.0', '0.10.1', '0.10.2', 'nightly']
+        jax-version: ['0.8.3', '0.9.0', '0.9.1', '0.9.2', '0.10.0', '0.10.1', '0.10.2', '0.11.0', 'nightly']
+        exclude:
+          - python-version: '3.11'
+            jax-version: '0.11.0'
 
     runs-on: ubuntu-latest
     timeout-minutes: 30
@@ -34,17 +40,21 @@ jobs:
       cancel-in-progress: true
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+      with:
+        persist-credentials: false
 
     # Install deps
-    - uses: actions/setup-python@v6
+    - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
       with:
         python-version: ${{ matrix.python-version }}
     - run: pip --version
-    - if: ${{ matrix.jax-version == 'nightly' }}
+    - if: matrix.jax-version == 'nightly'
       run: pip install -e . -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
-    - if: ${{ matrix.jax-version != 'nightly' }}
-      run: pip install -e . "jax==${{ matrix.jax-version }}"
+    - if: matrix.jax-version != 'nightly'
+      run: pip install -e . "jax==${JAX_VERSION}"
+      env:
+        JAX_VERSION: ${{ matrix.jax-version }}
     - run: pip freeze
 
     # Run tests