def test_custom_root_errors(self):
    with self.assertRaisesRegex(TypeError, re.escape('f() output pytree')):
        lax.custom_root((lambda x: (x, x)), 0.0, (lambda f, x: x), (lambda f, x: x))
    with self.assertRaisesRegex(TypeError, re.escape('solve() output pytree')):
        lax.custom_root((lambda x: x), 0.0, (lambda f, x: (x, x)), (lambda f, x: x))

    def dummy_root_usage(x):
        f = (lambda y: (x - y))
        return lax.custom_root(f, 0.0, (lambda f, x: x), (lambda f, x: (x, x)))
    with self.assertRaisesRegex(TypeError, re.escape('tangent_solve() output pytree')):
        api.jvp(dummy_root_usage, (0.0,), (0.0,))
def test_custom_root_errors(self):
    with self.assertRaisesRegex(TypeError, re.escape('f() output pytree')):
        lax.custom_root((lambda x: (x, x)), 0.0, (lambda f, x: x), (lambda f, x: x))
    with self.assertRaisesRegex(TypeError, re.escape('solve() output pytree')):
        lax.custom_root((lambda x: x), 0.0, (lambda f, x: (x, x)), (lambda f, x: x))

    def dummy_root_usage(x):
        f = (lambda y: (x - y))
        return lax.custom_root(f, 0.0, (lambda f, x: x), (lambda f, x: (x, x)))
    with self.assertRaisesRegex(TypeError, re.escape('tangent_solve() output pytree')):
        api.jvp(dummy_root_usage, (0.0,), (0.0,))