用pytorch 2.1 加速 numpy 代码

• 参考 https://pytorch.org/blog/compiling-numpy-code/

• 在mac M2 机器上, 快了50%， 但没有好几倍。可能和依赖libomp有关

brew install libomp
python test_np.py

• test_np.py 代码如下

import time
import numpy as np

def kmeans(X, means):
    return np.argmin(np.linalg.norm(X - means[:, None], axis=2), axis=0)

t1 = time.time()
npts = 10_000_000
X = np.repeat([[5, 5], [10, 10]], [npts, npts], axis=0)
X = X + np.random.randn(*X.shape)  # 2 distinct "blobs"
means = np.array([[5, 5], [10, 10]])
np_pred = kmeans(X, means)
t2 = time.time()
import torch
compiled_fn = torch.compile(kmeans)
t3 = time.time()
compiled_pred = compiled_fn(X, means)
t4 = time.time()
assert np.allclose(np_pred, compiled_pred)
print(t2 - t1 , '时间torch = ' ,  t4 - t3 )