New JAX build from Google: adopt changes for test-jax.sh

NVIDIA · Nov 26, 2024 · d9d6cd4 · d9d6cd4
1 parent 9a6256a
commit d9d6cd4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/.github/container/test-jax.sh b/.github/container/test-jax.sh
@@ -26,7 +26,7 @@ jax_source_dir() {
 
 query_tests() {
     cd `jax_source_dir`
-    python build/build.py --configure_only
+    python build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --configure_only
     bazel query tests/... 2>&1 | grep -F '//tests:'
     exit
 }
@@ -191,5 +191,5 @@ pip install matplotlib
 ## Run tests
 
 cd `jax_source_dir`
-python build/build.py --configure_only
+python build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --configure_only
 bazel test ${BAZEL_TARGET} ${TEST_TAG_FILTERS} ${COMMON_FLAGS} ${EXTRA_FLAGS}