def body_fun(state):
    (num, i, count) = state
    return (num, lax.add(i, 1), inner_loop(i, count))
def body_fun(state):
    (num, i, count) = state
    return (num, lax.add(i, 1), inner_loop(i, count))