Vectorizing Tensor Products from Python to Matlab

I am in the process of converting some code from Python into Matlab. I have code working that produces the same results, but I am wondering if there may be a way to vectorize some of my for loops in the Matlab code as it take a long time to run. X in an Nxd matrix, diff is an NxNxd tensor, kxy is an NxN matrix, gradK is an NxNx2 tensor, and sumkxy, dxkxy, and obj are all Nxd matrices.

Here is the original Python Code:

diff = x[:, None, :] - x[None, :, :]  # D_{ij, s}
kxy = np.exp(-np.sum(diff ** 2, axis=-1) / (2 * h ** 2)) / np.power(np.pi * 2.0 * h * h, d / 2) # -1 last dimension K_{ij]
gradK = -diff * kxy[:, :, None] / h ** 2  # N * N * 2
sumkxy = np.sum(kxy, axis=1)
dxkxy = np.sum(gradK, axis=1)  # N * 2    sum_{i} d_i K_{ij, s}
obj = np.sum(gradK / sumkxy[None, :, None], axis=1)  # N * 2

and here is my initial Matlab Code with all the for loops:

diff = zeros([n,n,d]);
for i = 1:n
    for j = 1:n
        for k = 1:d
            diff(i,j,k) = x(i,k) - x(j,k);
        end
    end
end

kxy = exp(-sum(dif.^2, 3)/(2*h^2))/((2*pi*h^2)^(d/2));
sumkxy = sum(kxy,2);
gradK = zeros([n,n,d]);
for i = 1:n
    for j = 1:n
        for k = 1:d
            gradK(i,j,k) = -diff(i,j,k)*kxy(i, j)/h^2;
        end
    end
end

dxkxy = squeeze(sum(gradK,2));
a = zeros([n,n,d]);
for i =1:n
    for j = 1:n
        for k = 1:d
            a(i,j,k) = gradK(i,j,k)/sumkxy(i);
        end
    end
end
obj = squeeze(sum(a, 2));

I know a way a faster way to calculate the kxy term is to use the following code:

XY = x*x';
x2= sum(x.^2, 2);
X2e = repmat(x2, 1, n);
 
H = (X2e + X2e' - 2*XY); % calculate pairwise distance
Kxy = exp(-H/(2*h^2))/((2*pi*h*h)^(d/2));

But then I struggle on a way to then calculate gradK efficiently without diff. Any help or suggestions would be greatly appreciated!

Answer

If your goal is computation of obj you don’t need even to compute gradK and a:

sx = sum(x.^2, 2);
H = sx - 2*x*x.' + sx.';

kxy = exp(-H/(2*h^2))/((2*pi*h^2)^(d/2));
kh = kxy / h^2; 
sumkxy = sum(kxy, 2);
khs = kh ./ sumkxy;
obj = khs * x - sum(khs, 2) .* x;

gradK and dif can be computed this way:

dif = reshape(x, n, 1, d) - reshape(x, 1, n, d);
gradK = -dif .* (kxy / h^2);.