Merge pull request #1314 from AI-Hypercomputer:parambole/dependency_update

PiperOrigin-RevId: 733170071

Author: maxtext authors

MaxText/tests/quantizations_test.py

@@ -99,7 +99,7 @@ def test_configure_quantization_is_int8(self):
       quant = _configure_quantization(quant_str="int8", mode_str=quant_mode)
       self.assertNotEqual(quant, None)
 
-  @pytest.mark.skip(reason="b/399968784 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
+  @pytest.mark.skip(reason="b/400476456 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
   def test_aqt_quantization(self):
     # Without quantization
     inputs, res_einsum, res_dg = _apply()

MaxText/tests/train_compile_test.py

@@ -304,7 +304,7 @@ def test_moe_dropping_bf16(self):
         )
     )
 
-  @pytest.mark.skip(reason="b/399968784 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
+  @pytest.mark.skip(reason="b/400476456 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
   @pytest.mark.tpu_only
   def test_moe_dropping_int8(self):
     compiled_trainstep_file = "/tmp/test_moe_dropping_int8.pickle"
@@ -391,7 +391,7 @@ def test_moe_dense_bf16(self):
         )
     )
 
-  @pytest.mark.skip(reason="b/399968784 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
+  @pytest.mark.skip(reason="b/400476456 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
   @pytest.mark.tpu_only
   def test_moe_dense_int8(self):
     compiled_trainstep_file = "/tmp/test_moe_dense_int8.pickle"

requirements.txt

@@ -29,7 +29,7 @@ sentencepiece==0.1.97
 tensorflow-text>=2.13.0
 tensorflow>=2.13.0
 tensorflow-datasets
-tensorboardx
+tensorboardx>=2.6.2.2
 tensorboard-plugin-profile
 tiktoken
 transformers

requirements_with_jax_stable_stack.txt

@@ -14,6 +14,7 @@ pytype
 sentencepiece==0.1.97
 tensorflow-text>=2.13.0
 tensorflow-datasets
+tensorboardx>=2.6.2.2
 tiktoken
 transformers
 mlperf-logging@git+https://github.com/mlperf/logging.git