From 5145c5de40858d0e8e5e4c3a4d6f92b2211157da Mon Sep 17 00:00:00 2001
From: Alexey Radul <axch@google.com>
Date: Mon, 17 Oct 2022 14:01:17 -0400
Subject: [PATCH] Better Jax installation instructions for benchmarking.

---
 .github/workflows/bench.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/bench.yaml b/.github/workflows/bench.yaml
index f986d4099..83e4f870e 100644
--- a/.github/workflows/bench.yaml
+++ b/.github/workflows/bench.yaml
@@ -29,7 +29,7 @@ jobs:
     - name: Install system dependencies
       run: |
         ${{ matrix.install_deps }}
-        pip install numpy jax
+        pip install numpy "jax[cpu]"
         echo "${{ matrix.path_extension }}" >> $GITHUB_PATH
 
     - name: Cache