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)