From 1817ee90a096f8c705e9c962684ad0d1aa38d918 Mon Sep 17 00:00:00 2001 From: Junpeng Lao Date: Wed, 19 Jul 2023 22:04:08 +0200 Subject: [PATCH] Update usage of soon-to-be deprecated Chex assertions (#560) --- tests/mcmc/test_integrators.py | 2 +- tests/mcmc/test_latent_gaussian.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/mcmc/test_integrators.py b/tests/mcmc/test_integrators.py index 3e300d839..68f1dbd88 100644 --- 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 index 07143f133..32fe6cc12 100644 --- 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 )