I have three tensors X,Y,Z of dimensions (K,) (K,) (n,) and a function k that takes two inputs and returns a float.

I would like to convert the following loop to obtain a tensor of dimensions (K,K,n)

M = torch.zeros( (len(X),len(Y),len(Z)))

for i in range(len(X)):
    for j in range(len(Y)):
        for s in range(len(Z)):
            M[i,j,s] = k(X[i], Z[s]) * k(Y[i], Z[s])

I cannot manage to get the right dimensions to match using broadcasting.