def loop(init):
    result = lax.while_loop(loop_cond, loop_body, (init, 0))
    (_, count) = result
    return count
def loop(init):
    result = lax.while_loop(loop_cond, loop_body, (init, 0))
    (_, count) = result
    return count