def outer_loop(arr):

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

    def body_fun(state):
        (i, num, arr, out) = state
        return (lax.add(i, 1), num, arr, inner_loop(i, arr, out))
    out = np.zeros(arr.shape, dtype=arr.dtype)
    init_val = (0, num, arr, out)
    (_, _, _, out) = lax.while_loop(cond_fun, body_fun, init_val)
    return out
def outer_loop(arr):

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

    def body_fun(state):
        (i, num, arr, out) = state
        return (lax.add(i, 1), num, arr, inner_loop(i, arr, out))
    out = np.zeros(arr.shape, dtype=arr.dtype)
    init_val = (0, num, arr, out)
    (_, _, _, out) = lax.while_loop(cond_fun, body_fun, init_val)
    return out