在pytorch（或Numpy）中实现这个等式的更有效方法

Question

我正在实现这个函数的分析形式

其中k（x，y）是RBF核k(x,y) = exp(-||x-y||^2 / (2h))

我的功能原型是

def A(X, Y, grad_log_px，Kxy):
   pass

而X，Y是NxD矩阵，其中N是批量大小，D是一个维度。所以X是一批x，在上面的等式中，N是我用autograd计算的NxD矩阵。在

Kxy是NxN矩阵，其中每个条目(i,j)是RBF核K(X[i],Y[j])

这里的挑战在于，在上面的等式中，y只是一个维数为D的向量。我有点想传入y的一批。（所以要传递Y大小的矩阵NxD）

这个公式可以通过批量大小循环使用，但是我很难以更简洁的方式实现

下面是我尝试的循环解决方案：

^{pr2}$

在方程中：grad_{x}和{}两个维度D

Answer 1

考虑到我正确地推断了各种术语的所有维度，这里有一个方法。但首先是维度的摘要（屏幕截图，因为使用数学类型设置更容易解释；请验证它们是否正确）：

另请注意第二项的双导数，它给出：

其中下标表示样本，上标表示特征。在

所以我们可以通过使用^{}（类似地^{}）和array broadcasting来创建这两个术语：

grad_y_K = (X[:, None, :] - Y) / h * K[:, :, None]  # Shape: N_x, N_y, D
term_1 = np.einsum('ij,ikl->ikjl', grad_log_px, grad_y_K)  # Shape: N_x, N_y, D_x, D_y
term_2_h = np.einsum('ij,kl->ijkl', K, np.eye(D)) / h  # Shape: N_x, N_y, D_x, D_y
term_2_h2_xy = np.einsum('ijk,ijl->ijkl', grad_y_K, grad_y_K)  # Shape: N_x, N_y, D_x, D_y
term_2_h2 = K[:, :, None, None] * term_2_h2_xy / h**2  # Shape: N_x, N_y, D_x, D_y
term_2 = term_2_h - term_2_h2  # Shape: N_x, N_y, D_x, D_y

结果如下：

^{pr2}$