def loop_body(state):
    effect[0] = True
    (pos, count) = state
    f = (lambda pos, inc: (lax.add(pos, 1), lax.add(count, inc)))
    return api.jit(f)(pos, inc)
def loop_body(state):
    effect[0] = True
    (pos, count) = state
    f = (lambda pos, inc: (lax.add(pos, 1), lax.add(count, inc)))
    return api.jit(f)(pos, inc)