Update test-requirements.txt (#155)

* Update test-requirements.txt * Update test-requirements.txt * run builddeps
EnzymeAD · Nov 7, 2024 · 4e89540 · 4e89540
1 parent 140da3b
commit 4e89540
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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'