使用 Numpy einsum 在 Python 中向量化函数

Question

I'm trying to use Numpy to vectorize a function for calculating the average marginal effects of the features in a multinomial logit model fitted using scikit-learn. I have managed to do the calculations using for loops , which looks like this

#Get probabilities for each obs i belonging to class j. shape =  N * J
probas = fitted_model.predict_proba(X)
#Get coefficients. Shape j_classes * k_coefficients
betas = fitted_model.coef_

J = probas.shape[1]
N = probas.shape[0]
K = betas.shape[1]

avg_margins = np.zeros([K, J])
        
for j in tqdm(range(J)):
    for k in range(K):
        dydw = 0
        for i in range(N):
            dydw += probas[i,j] * (betas[j,k] - np.dot(probas[i,:],betas[:,k]))
        avg_margins[k,j] = 1 / N * dydw

However, this is very slow, why I want someway of getting rid of the loops. I'm a beginner with both Numpy and linear algebra so please bare with me, but I think my best bet is using numpy einsum and where I've gotten so far is this

avg_margins = 1 / N * np.einsum('ij, jk -> kj', probas,  betas - np.einsum('im, mk -> k', probas, betas))

This unfortunately returns the wrong results and I'm unsure where I'm going wrong. Any help or hints of where I'm going wrong would be highly appreciated!

Answer 1

No einsum is needed here:

diff = betas[:,None] - np.dot(probas, betas)
avg_margins = np.sum(probas * diff.T, axis=1) / N