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 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