Update submodules, prepare for leasing v0.2.4

Author: qihqi

benchmarks/prefill_offline.py

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

benchmarks/run_offline.py

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

deps/JetStream

@@ -25,7 +25,7 @@ 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 jax[tpu]==0.4.29 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 pip install tensorflow-text
 pip install tensorflow
 

deps/xla

@@ -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