def test_custom_root_scalar(self):

    def scalar_solve(f, y):
        return (y / f(1.0))

    def binary_search(func, x0, low=0.0, high=100.0):
        del x0

        def cond(state):
            (low, high) = state
            midpoint = (0.5 * (low + high))
            return ((low < midpoint) & (midpoint < high))

        def body(state):
            (low, high) = state
            midpoint = (0.5 * (low + high))
            update_upper = (func(midpoint) > 0)
            low = jnp.where(update_upper, low, midpoint)
            high = jnp.where(update_upper, midpoint, high)
            return (low, high)
        (solution, _) = lax.while_loop(cond, body, (low, high))
        return solution

    def sqrt_cubed(x, tangent_solve=scalar_solve):
        f = (lambda y: ((y ** 2) - (x ** 3)))
        return lax.custom_root(f, 0.0, binary_search, tangent_solve)
    (value, grad) = api.value_and_grad(sqrt_cubed)(5.0)
    self.assertAllClose(value, (5 ** 1.5), check_dtypes=False, rtol=1e-06)
    self.assertAllClose(grad, api.grad(pow)(5.0, 1.5), check_dtypes=False, rtol=1e-07)
    jtu.check_grads(sqrt_cubed, (5.0,), order=2, rtol={jnp.float32: 0.01, jnp.float64: 0.001})
    inputs = jnp.array([4.0, 5.0])
    results = api.vmap(sqrt_cubed)(inputs)
    self.assertAllClose(results, (inputs ** 1.5), check_dtypes=False)
    results = api.jit(sqrt_cubed)(5.0)
    self.assertAllClose(results, (5.0 ** 1.5), check_dtypes=False, rtol={np.float64: 1e-07})
def test_custom_root_scalar(self):

    def scalar_solve(f, y):
        return (y / f(1.0))

    def binary_search(func, x0, low=0.0, high=100.0):
        del x0

        def cond(state):
            (low, high) = state
            midpoint = (0.5 * (low + high))
            return ((low < midpoint) & (midpoint < high))

        def body(state):
            (low, high) = state
            midpoint = (0.5 * (low + high))
            update_upper = (func(midpoint) > 0)
            low = jnp.where(update_upper, low, midpoint)
            high = jnp.where(update_upper, midpoint, high)
            return (low, high)
        (solution, _) = lax.while_loop(cond, body, (low, high))
        return solution

    def sqrt_cubed(x, tangent_solve=scalar_solve):
        f = (lambda y: ((y ** 2) - (x ** 3)))
        return lax.custom_root(f, 0.0, binary_search, tangent_solve)
    (value, grad) = api.value_and_grad(sqrt_cubed)(5.0)
    self.assertAllClose(value, (5 ** 1.5), check_dtypes=False, rtol=1e-06)
    self.assertAllClose(grad, api.grad(pow)(5.0, 1.5), check_dtypes=False, rtol=1e-07)
    jtu.check_grads(sqrt_cubed, (5.0,), order=2, rtol={jnp.float32: 0.01, jnp.float64: 0.001})
    inputs = jnp.array([4.0, 5.0])
    results = api.vmap(sqrt_cubed)(inputs)
    self.assertAllClose(results, (inputs ** 1.5), check_dtypes=False)
    results = api.jit(sqrt_cubed)(5.0)
    self.assertAllClose(results, (5.0 ** 1.5), check_dtypes=False, rtol={np.float64: 1e-07})