[Tests] Adjust TPU test values (#1233)

* [Tests] Adjust TPU test values * slow tests * remaining refs

[Tests] Adjust TPU test values (#1233)
* [Tests] Adjust TPU test values * slow tests * remaining refs
2e980ac9 · Anton Lozhkov · GitHub · 0feb21a1 · 2e980ac9
Unverified Commit 2e980ac9 authored Nov 10, 2022 by Anton Lozhkov Committed by GitHub Nov 10, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 5 deletions

tests/test_scheduler_flax.py tests/test_scheduler_flax.py +5 -5

No files found.
--- a/tests/test_scheduler_flax.py
+++ b/tests/test_scheduler_flax.py
@@ -859,7 +859,7 @@ class FlaxPNDMSchedulerTest(FlaxSchedulerCommonTest):
        result_mean = jnp.mean(jnp.abs(sample))

        if jax_device == "tpu":
-            assert abs(result_sum - 198.1542) < 1e-2
+            assert abs(result_sum - 198.1275) < 1e-2
            assert abs(result_mean - 0.2580) < 1e-3
        else:
            assert abs(result_sum - 198.1318) < 1e-2
@@ -872,8 +872,8 @@ class FlaxPNDMSchedulerTest(FlaxSchedulerCommonTest):
        result_mean = jnp.mean(jnp.abs(sample))

        if jax_device == "tpu":
-            assert abs(result_sum - 185.4352) < 1e-2
-            assert abs(result_mean - 0.24145) < 1e-3
+            assert abs(result_sum - 186.83226) < 1e-2
+            assert abs(result_mean - 0.24327) < 1e-3
        else:
            assert abs(result_sum - 186.9466) < 1e-2
            assert abs(result_mean - 0.24342) < 1e-3
@@ -885,8 +885,8 @@ class FlaxPNDMSchedulerTest(FlaxSchedulerCommonTest):
        result_mean = jnp.mean(jnp.abs(sample))

        if jax_device == "tpu":
-            assert abs(result_sum - 185.4352) < 1e-2
-            assert abs(result_mean - 0.2414) < 1e-3
+            assert abs(result_sum - 186.83226) < 1e-2
+            assert abs(result_mean - 0.24327) < 1e-3
        else:
            assert abs(result_sum - 186.9482) < 1e-2
            assert abs(result_mean - 0.2434) < 1e-3