More efficient way of implement this equation in pytorch (or Numpy)
I'm implementing the analytical form of this function



where k(x,y) is a RBF kernel k(x,y) = exp(-||x-y||^2 / (2h))

My function prototype is 

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


and X, Y are NxD matrix where N is batch size and D is a dimension. So X is a batch of x with size N in the above equation grad_log_px is some NxD matrix I've computed using autograd.

Kxy is NxN matrix where each entry (i,j) is the RBF kernel K(X[i],Y[j])

The challenge here is that in the above equation, y is just a vector with dimension D. I kind of want to pass into a batch of y. (So to pass matrix Y with NxD size)

The equation is fine using loop through the batch size but I'm having trouble to implement in a more neat way

here is my attempted loop solution:

def A(X, Y, grad_log_px，Kxy):
   res = []
   for i in range(Y.shape[0]):
       temp = 0
       for j in range(X.shape[0]):
           # first term of equation
           temp += grad_log_px[j].reshape(D,1)@(Kxy[j,i] * (X[i] - Y[j]) / h).reshape(1,D)
           temp += Kxy[j,i] * np.identity(D) - ((X[i] - Y[j]) / h).reshape(D,1)@(Kxy[j,i] * (X[i] - Y[j]) / h).reshape(1,D) # second term of equation
       temp /= X.shape[0]

        res.append(temp)
    return np.asarray(res) # return NxDxD array 


In the equation: grad_{x} and grad_{y} both dimension D

Given that that I inferred all the dimensions of the various terms correctly here's a way to go about it. But first a summary of the dimensions (screenshot as it's easier to explain with math type setting; please verify if they are correct):



Also note the double derivative of the second term which gives:



where subscripts denote samples and superscripts denote features.

So we can create the two terms by using np.einsum (similarly torch.einsum) and array broadcasting:

grad_y_K = (X[:, None, :] - Y) / h * K[:, :, None]  # Shape: N_x, N_y, D
term_1 = np.einsum('ij,ikl-&gt;ikjl', grad_log_px, grad_y_K)  # Shape: N_x, N_y, D_x, D_y
term_2_h = np.einsum('ij,kl-&gt;ijkl', K, np.eye(D)) / h  # Shape: N_x, N_y, D_x, D_y
term_2_h2_xy = np.einsum('ijk,ijl-&gt;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


Then the result is given by:

(term_1 + term_2).sum(axis=0) / N  # Shape: N_y, D_x, D_y

