Повышение эффективности алгоритма начальной загрузки для большой матрицы

Question

Я работаю над повышением эффективности алгоритма начальной загрузки, написанного на python. Скорость вычислений в порядке, когда размер матрицы относительно мал, но она значительно замедляется, когда n и k большие.

Поскольку алгоритм по сути является манипулированием массивами, можно ли в любом случае повысить его эффективность, либо оптимизировав сам алгоритм, либо переписав его, чтобы использовать такие пакеты, как pycuda / tenorflow и т. Д., Поскольку ресурсы GPU доступны?

import numpy as np
n = 5000
k = 3000
ary = np.random.choice([0,1],size=(n,k))
lst = []
cols = list(range(ary.shape[1]))
while len(lst) < ary.shape[1]:
    print(len(lst), ary.shape[1])
    a = [lst+[x] for x in cols]
    b = [ary[:,x] for x in a]
    c = np.asarray(b)
    d = c.sum(axis=2)
    d = np.repeat(d[..., None], c.shape[-1], 2)
    d1 = np.divide(c, d)
    d2 = np.nanmean(d1, axis=1)
    d2 = np.nan_to_num(d2)
    d3 = d2[:, -1]
    d4 = d3 / np.sum(d3)
    lst.extend([np.random.choice(cols, p=d4)])
    print(d4.shape)
    print(lst)