def inner_loop(i, count):

    def cond_fun(state):
        (i, j, _) = state
        return lax.le(j, i)

    def body_fun(state):
        (i, j, count) = state
        return (i, lax.add(j, 1), lax.add(count, 1))
    init_val = (i, 0, count)
    (_, _, count) = lax.while_loop(cond_fun, body_fun, init_val)
    return count
def inner_loop(i, count):

    def cond_fun(state):
        (i, j, _) = state
        return lax.le(j, i)

    def body_fun(state):
        (i, j, count) = state
        return (i, lax.add(j, 1), lax.add(count, 1))
    init_val = (i, 0, count)
    (_, _, count) = lax.while_loop(cond_fun, body_fun, init_val)
    return count