Fast Kendall Tau calculation with pytorch.

kendall_tau_pytorch.py

import torch
import time
from scipy.stats import kendalltau

def kendall(x, y):
    n = x.shape[0]
    
    def sub_pairs(x):
        return x.expand(n,n).T.sub(x).sign_()
    
    return sub_pairs(x).mul_(sub_pairs(y)).sum().div(n*(n-1))

d = torch.empty(10)
for i in range(10):
    x, y = torch.randperm(4000), torch.randperm(4000)
    t = time.time_ns()
    m = kendall(x,y)
    d[i] = (time.time_ns() - t)*1e-6
    print(f'{d[i]:.2f}ms')
    print(f'{abs(kendalltau(x,y).correlation - m):.9f}')
print(f'AVG {d.mean():.2f}ms')