def testFindByValueNotFound(self, value): def f(x): y = x**2 z = 1 - y return z find_by_value(f, value)(jnp.array([1.0, 2.0, 3.0]))
def testFindByValueFound(self, value): def f(x): y = x**2 z = 1 - y r = 1 / z return r * 0 with self.assertRaises(FoundValue): find_by_value(f, value)(jnp.array([1.0, 2.0, 3.0]))