Merge pull request #1069 from carlosgmartin:relax_atol_assert_trees_a…

…ll_close

PiperOrigin-RevId: 678616122

optax/_src/alias_test.py

@@ -587,7 +587,9 @@ def test_plain_preconditioning(self):
     expected_precond_vec = precond_mat.dot(
         vec, precision=jax.lax.Precision.HIGHEST
     )
-    chex.assert_trees_all_close(plain_precond_vec, expected_precond_vec)
+    chex.assert_trees_all_close(
+        plain_precond_vec, expected_precond_vec, rtol=1e-5
+    )
 
   @parameterized.product(idx=[0, 1, 2, 3])
   def test_preconditioning_by_lbfgs_on_vectors(self, idx: int):

optax/_src/linesearch_test.py

@@ -485,8 +485,8 @@ def test_linesearch(self, problem_name: str, seed: int):
     with self.subTest('Check against scipy'):
       stepsize = otu.tree_get(final_state, 'learning_rate')
       final_value = otu.tree_get(final_state, 'value')
-      chex.assert_trees_all_close(scipy_res[0], stepsize, atol=1e-5)
-      chex.assert_trees_all_close(scipy_res[3], final_value, atol=1e-5)
+      chex.assert_trees_all_close(scipy_res[0], stepsize, rtol=1e-5)
+      chex.assert_trees_all_close(scipy_res[3], final_value, rtol=1e-5)
 
   def test_failure_descent_direction(self):
     """Check failure when updates are not a descent direction."""

optax/contrib/_dog.py

@@ -175,11 +175,11 @@ def dog(
     ...    grad, opt_state, params, value=value)
     ...  params = optax.apply_updates(params, updates)
     ...  print('Objective function: ', f(params))
-    Objective function:  13.999964
-    Objective function:  13.999941
-    Objective function:  13.999905
-    Objective function:  13.999857
-    Objective function:  13.999794
+    Objective function:  13.99...
+    Objective function:  13.99...
+    Objective function:  13.99...
+    Objective function:  13.99...
+    Objective function:  13.99...
 
   References:
     Ivgi et al., `DoG is SGD's Best Friend: A Parameter-Free Dynamic Step