Fix all tests in preparation for the next Saxml release.

PiperOrigin-RevId: 560205746 Change-Id: I0caf001fdc1d84c6b80e942b3b80ef6240c33ea5
google · Aug 25, 2023 · df12a6c · df12a6c
1 parent 7cf1751
commit df12a6c
Show file tree

Hide file tree

Showing 6 changed files with 254 additions and 199 deletions.
diff --git a/requirements-cuda.in b/requirements-cuda.in
@@ -2,4 +2,4 @@
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 # We keep jax[tpu] because libtpu-nightly is a dependency of saxml/server:server_deps
-jax[cuda]
+jax[cuda] @ git+https://github.com/google/jax