def test_tree_all_finite_should_fail_inf(self):
inf_tree = {
'finite_var': jnp.ones((3,)),
'inf_var': jnp.array([0.0, jnp.inf]),
}
with self.assertRaisesRegex(AssertionError,
'Tree contains non-finite value'):
asserts.assert_tree_all_finite(inf_tree)
def test_tree_all_finite_passes_finite(self):
finite_tree = {'a': jnp.ones((3,)), 'b': jnp.array([0.0, 0.0])}
asserts.assert_tree_all_finite(finite_tree)