Задан массив ковариационной матрицы; Как эффективно преобразовать каждый из них в матрицу корреляции?

Question

Предположим, у меня есть массив матриц M.shape = (N,D,D), где D - размерность матрицы ковариации, а N - количество матриц ковариации.

У меня есть метод преобразования одиночная ковариационная матрица к корреляционной матрице, определенной как

def convert_to_correlation(m):
    inv = torch.inverse(torch.diag(torch.diag(s)).sqrt())
    r = inv@s@inv
    return r

Когда у меня есть список такой ковариационной матрицы, я работаю как

def convert_all_to_correlation(M):
    res = []
    for i in range(len(M)):
        m = convert_to_correlation(M[i])
        res.append(m)
    res = torch.stack(res)
    return res

Есть две проблемы выше подходить;

(1) Массив может иметь большое первое измерение; следовательно, в итеративном цикле очень медленно (2) My M является выходом нейронной сети (следовательно, является частью вычислительного графа); Я не хочу, чтобы мой подход convert_to_all_correlation удалил их из графика.

Любая помощь приветствуется! Спасибо