Add JAX requirements

2026-03-25 09:59:12 +08:00 · 2024-04-16 08:05:54 +00:00 · 2024-04-16 08:05:54 +00:00 · eb0a0466a9
commit eb0a0466a9
parent c59c1e7b2c
1 changed files with 3 additions and 0 deletions
--- a/requirements-tpu.txt
+++ b/requirements-tpu.txt
@ -1,2 +1,5 @@
 # Common dependencies
 -r requirements-common.txt
+
+jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+flax >= 0.8