def _while_loop(cond_fun, body_fun, init_val, max_iter):

    def _iter(val):
        next_val = body_fun(val)
        next_cond = True
        return (next_val, next_cond)

    def _fun(tup, _):
        (val, cond) = tup
        return (jax.lax.cond(cond, _iter, (lambda x: (x, False)), val), _)
    init = (init_val, cond_fun(init_val))
    return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]
def _while_loop(cond_fun, body_fun, init_val, max_iter):

    def _iter(val):
        next_val = body_fun(val)
        next_cond = True
        return (next_val, next_cond)

    def _fun(tup, _):
        (val, cond) = tup
        return (jax.lax.cond(cond, _iter, (lambda x: (x, False)), val), _)
    init = (init_val, cond_fun(init_val))
    return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]