Computes the covariance matrix in PyTorch.

cov.py

"""Computes the covariance matrix in PyTorch."""


def cov_mean(input, unbiased=True, keepdims=False):
    n = input.shape[-1] - unbiased
    mean = input.mean(dim=-1, keepdims=True)
    dev = input - mean
    mean = mean if keepdims else mean[..., 0]
    return dev @ dev.transpose(-1, -2) / n, mean


def cov(input, unbiased=True):
    return cov_mean(input, unbiased=unbiased)[0]