Update submodules, prepare for leasing v0.2.4

benchmarks/prefill_offline.py

@@ -16,6 +16,9 @@
 import os
 import time
 
+# import torch_xla2 first!
+# pylint: disable-next=all
+import torch_xla2
 import humanize
 import jax
 import numpy as np

benchmarks/run_offline.py

@@ -16,6 +16,9 @@
 import os
 import time
 
+# import torch_xla2 first!
+# pylint: disable-next=all
+import torch_xla2
 import jax
 import jax.numpy as jnp
 # pylint: disable-next=all

install_everything.sh

@@ -24,18 +24,18 @@ pip show tensorboard && pip uninstall -y tensorboard
 pip show tensorflow-text && pip uninstall -y tensorflow-text
 pip show torch_xla2 && pip uninstall -y torch_xla2
 
-pip install flax==0.8.3
-pip install jax[tpu]==0.4.28 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+pip install flax
 pip install tensorflow-text
 pip install tensorflow
 
 pip install ray[default]==2.22.0
 # torch cpu
-pip install torch==2.2.1+cpu --index-url https://download.pytorch.org/whl/cpu
+pip install torch==2.3.1+cpu --index-url https://download.pytorch.org/whl/cpu
 pip install tensorflow flatbuffers absl-py sentencepiece seqio google-cloud-storage 
 pip install safetensors colorama coverage humanize
 
 git submodule update --init --recursive
 pip show google-jetstream && pip uninstall -y google-jetstream
 pip show torch_xla2 && pip uninstall -y torch_xla2
 pip install -e .
+pip install -U jax[tpu]==0.4.29 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

run_interactive.py

@@ -17,6 +17,8 @@
 import time
 from typing import List
 
+# import torch_xla2 first!
+import torch_xla2  # pylint: disable
 import jax
 import numpy as np
 from absl import app, flags

run_server_with_ray.py

@@ -18,6 +18,8 @@
 from typing import Sequence
 from absl import app, flags
 
+# import torch_xla2 first!
+import torch_xla2  # pylint: disable
 import jax
 from jetstream.core import server_lib
 from jetstream.core.config_lib import ServerConfig