def loop(init, local_limit, inc):

    def loop_cond(state):
        (pos, _) = state
        return lax.lt(pos, local_limit)

    def loop_body(state):
        effect[0] = True
        (pos, count) = state
        return (lax.add(pos, 1), lax.add(count, inc))
    result = lax.while_loop(loop_cond, loop_body, (init, 0))
    (_, count) = result
    return count
def loop(init, local_limit, inc):

    def loop_cond(state):
        (pos, _) = state
        return lax.lt(pos, local_limit)

    def loop_body(state):
        effect[0] = True
        (pos, count) = state
        return (lax.add(pos, 1), lax.add(count, inc))
    result = lax.while_loop(loop_cond, loop_body, (init, 0))
    (_, count) = result
    return count