Updated scheduler test values

Author: eltsai

tests/schedulers/test_scheduler_flax.py

@@ -335,8 +335,8 @@ def test_full_loop_no_noise(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 263.11) < 1.5e-2
-            assert abs(result_mean - 0.34259) < 2e-5
+            assert abs(result_sum - 257.32495) < 1.5e-2
+            assert abs(result_mean - 0.335059) < 2e-5
         else:
             assert abs(result_sum - 255.1113) < 1e-2
             assert abs(result_mean - 0.332176) < 1e-3
@@ -621,7 +621,7 @@ def test_full_loop_with_set_alpha_to_one(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 149.8409) < 1e-2
+            assert abs(result_sum - 149.82944) < 1e-2
             assert abs(result_mean - 0.1951) < 1e-3
         else:
             assert abs(result_sum - 149.8295) < 1e-2
@@ -919,7 +919,7 @@ def test_full_loop_with_set_alpha_to_one(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 186.83226) < 8e-2
+            assert abs(result_sum - 186.94574) < 8e-2
             assert abs(result_mean - 0.24327) < 1e-3
         else:
             assert abs(result_sum - 186.9466) < 1e-2
@@ -932,7 +932,7 @@ def test_full_loop_with_no_set_alpha_to_one(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 186.83226) < 8e-2
+            assert abs(result_sum - 186.94574) < 8e-2
             assert abs(result_mean - 0.24327) < 1e-3
         else:
             assert abs(result_sum - 186.9482) < 1e-2