Cheap way to turn covariance matrix in tensorflow to correlation matrix-CodePudding

I have tensor of shape [N,d,d] of N covariance matrices and I would like to turn them into correlation matrices. Is there any cheap implementation in tensorflow / keras?

CodePudding user response：

You already have the correlation matrix, so all you need to do is divide each element of that matrix by the appropriate standard deviation product. These standard deviations are present on your matrix's diagonals, and can be efficiently retrieved with a few matrix multiplications (more about the math on this wikipedia page):

import tensorflow as tf

def to_corr(cov):
    d_inv = tf.linalg.diag(1/tf.sqrt(tf.abs(tf.linalg.diag_part(cov))))
    corr = d_inv @ cov @ d_inv
    return corr

Sanity check:

import tensorflow_probability as tfp

x = tf.random.normal(shape=(3,5,3))
cov_matrix = tfp.stats.covariance(x, sample_axis=0, event_axis=-1)
tf.debugging.assert_near(to_corr(cov_matrix),tfp.stats.correlation(x, sample_axis=0, event_axis=-1))