def count(num):

    def body_fun(i, tot):
        return lax.add(num, lax.add(tot, i))
    return lax.fori_loop(0, num, body_fun, 0)
def count(num):

    def body_fun(i, tot):
        return lax.add(num, lax.add(tot, i))
    return lax.fori_loop(0, num, body_fun, 0)