diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 77a435b7..3e2c57fe 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -47,6 +47,7 @@ steps:
         mv bazel* .local/bin/bazel.exe
       fi
       mkdir -p .baztmp
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp run //builddeps:requirements.update
       HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --test_output=errors //test/...
       HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --cache_test_results=no //test:bench_vs_xla
       HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --cache_test_results=no //test:llama
diff --git a/builddeps/test-requirements.txt b/builddeps/test-requirements.txt
index 1de2167a..347925f1 100644
--- a/builddeps/test-requirements.txt
+++ b/builddeps/test-requirements.txt
@@ -1,11 +1,11 @@
 absl-py
-jax
 numpy
+jax
 jaxlib
 https://github.com/wsmoses/jax-md/archive/45059b8f63dad0b5cb171feafff71b82162487e7.tar.gz
 # maxtext can't be installed concurrently, but installing it fixes
 # https://github.com/wsmoses/maxtext/archive/bc50722be7d89e4003bd830b80e4ac968be658eb.tar.gz
-jax[cuda12_pip]; sys_platform == 'linux'
+jax-cuda12-plugin[with_cuda]; sys_platform == 'linux'
 requests; sys_platform == 'linux'
 # -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 # libtpu-nightly == 0.1.dev20240729; sys_platform == 'linux'