Update usage of soon-to-be deprecated Chex assertions (#560)

blackjax-devs · Jul 19, 2023 · 1817ee9 · 1817ee9
1 parent cf94b27
commit 1817ee9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/mcmc/test_integrators.py b/tests/mcmc/test_integrators.py
@@ -128,7 +128,7 @@ def test_integrator(self, example_name, integrator_name):
         )
 
         # We make sure that the particle moved from its initial position.
-        chex.assert_tree_all_close(final_state.position, example["q_final"], atol=1e-2)
+        chex.assert_trees_all_close(final_state.position, example["q_final"], atol=1e-2)
 
         # We now check the conservation of energy, the property that matters the most in HMC.
         energy = -neg_potential(q) + kinetic_energy(p)

diff --git a/tests/mcmc/test_latent_gaussian.py b/tests/mcmc/test_latent_gaussian.py
@@ -51,10 +51,10 @@ def body(carry, key):
                 np.zeros((D,)), C, obs, R
             )
 
-        chex.assert_tree_all_close(
+        chex.assert_trees_all_close(
             np.mean(states.position, 0), expected_mean, atol=1e-1, rtol=1e-2
         )
-        chex.assert_tree_all_close(
+        chex.assert_trees_all_close(
             np.cov(states.position, rowvar=False), expected_cov, atol=1e-1, rtol=1e-1
         )