def outer_loop(num):

    def cond_fun(state):
        (num, i, _) = state
        return lax.lt(i, num)

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

    def cond_fun(state):
        (num, i, _) = state
        return lax.lt(i, num)

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