update tests for jax 0.5.3

Author: matthew-e-hopkins

axlearn/audio/decoder_asr_test.py

@@ -569,7 +569,7 @@ def test_forward_summary(self):
         summaries = output_collections.summaries
         # 6 out of 8 examples are valid, therefore the average example weight is 0.75
         self._check_summary(summaries, "loss/example_weight", WeightedScalar(0.75, 8))
-        self._check_summary(summaries, "loss/ctc_loss", WeightedScalar(6972.1353, 6))
+        self._check_summary(summaries, "loss/ctc_loss", WeightedScalar(6972.135, 6))
         self._check_summary(summaries, "loss/invalid_seq_percent", 0.25)
         total_ctc_loss = summaries["loss/ctc_loss"].weight * summaries["loss/ctc_loss"].mean
         num_valid_frames = jnp.sum(safe_not(paddings) * per_example_weight[:, None])

axlearn/audio/frontend_test.py

@@ -236,6 +236,7 @@ def test_small_input(self, input_dtype):
                 output_with_large_mel_floor,
             ),
             output_with_correct_mel_floor,
+            rtol=6e-4,
         )
 
     @set_threefry_partitionable(False)  # TODO(Luzy): update for threefry_partitionable True
@@ -279,7 +280,7 @@ def test_fft(self):
             ref_outputs = self._jit_forward(ref_layer, inputs, paddings)
             test_outputs = self._jit_forward(layer, inputs, paddings)
 
-        self.assertAllClose(ref_outputs["outputs"], test_outputs["outputs"])
+        self.assertAllClose(ref_outputs["outputs"], test_outputs["outputs"], rtol=5e-3)
         self.assertAllClose(ref_outputs["paddings"], test_outputs["paddings"])
 
     @parameterized.product(

axlearn/audio/frontend_utils_test.py

@@ -417,7 +417,7 @@ def test_fft(self):
             ref_ffts = jax.jit(jnp.fft.fft, static_argnames="n")(inputs, n=fft_size)
             test_ffts = fft_fn(inputs)
 
-        assert_allclose(ref_ffts, test_ffts)
+        assert_allclose(ref_ffts, test_ffts, rtol=1e-3)
         # Run the following on gpu.
         jax.debug.inspect_array_sharding(test_ffts, callback=print)
         jax.debug.inspect_array_sharding(ref_ffts, callback=print)

axlearn/common/base_layer_test.py

@@ -41,7 +41,7 @@
     FanAxes,
     WeightInitializer,
 )
-from axlearn.common.test_utils import TestCase, assert_allclose, set_threefry_partitionable
+from axlearn.common.test_utils import TestCase, assert_allclose
 from axlearn.common.utils import safe_not
 
 
@@ -375,7 +375,6 @@ def test_apply_parameter_noise_recursively(self, param_noise_cfg):
                 self.assertNestedAllClose(jnp.zeros_like(orig_value), noisy_value)
 
     @parameterized.parameters(False, True)
-    @set_threefry_partitionable(True)  # TODO(mhopkins): remove during jax 0.5.0+ upgrade
     def test_tensor_stats(self, inline_child_summaries: bool):
         test_layer: TestLayer = (
             TestLayer.default_config()

axlearn/common/quantizer_test.py

@@ -348,7 +348,7 @@ def test_forward(
             ),
             expected_values[batch_size][normalize_codebook]["q_vecs"],
             atol=1e-6,
-            rtol=1e-6,
+            rtol=2e-6,
         )
         assert_allclose(
             np.sum(q_outputs.ids * safe_not(paddings)[:, :, None]),