如果我只是看看你发布的代码，我建议你试着删除内部循环，并使用Numpy的自动广播来做同样的事情。我注意到你正在以一种非常手动的方式计算欧氏距离，但distance module from SciPy似乎不支持autograd。
这就给我留下了下面的代码：

def potential(x):
    A, XI, P, Q, R0 = parameters()

    n_atoms = len(A)
    di = np.diag_indices(n_atoms)
    x = x.reshape(n_atoms, 3)

    # preprocess params to minimize work and remove need to special case "i == j"
    A[di] = 0.
    XI[di] = 0.
    P *= -1.
    Q *= -1.
    R0[di]= 1.

    U = 0.
    for i, xi in enumerate(x):
        dist = np.linalg.norm(xi - x, axis=1) / R0[i] - 1.
        Ub = XI[i] * np.exp(Q[i] * dist)
        Ur = A[i] * np.exp(P[i] * dist)
        U += sum(Ur) - sum(Ub)

    return U

当我查看你的上游代码库时，我看到你从optimize调用这个方法，这表明你在做数值优化的时候从文件中加载数据，这感觉非常错误。
因此，我建议进行重构以消除这种重复加载数据的情况，并清理参数。将内容存储在对象中似乎是有意义的，只需让它暴露potential和渐变的方法。

from autograd import elementwise_grad as egrad
import autograd.numpy as np

class GuptaPotential:
    def __init__(self):
        self.A, self.XI, self.P, self.Q, self.R0 = parameters()
        
        self.n_atoms = len(self.A)
        di = np.diag_indices(self.n_atoms)
        
        # preprocess params to minimize work and remove need to special case "i == j"
        self.A[di] = 0.
        self.XI[di] = 0.
        self.P *= -1
        self.Q *= -1
        self.R0[di] = 1.

    def potential(self, x):
        x = x.reshape(self.n_atoms, 3)
        U = 0.
        for i, xi in enumerate(x):
            dist = np.linalg.norm(xi - x, axis=1) / self.R0[i] - 1.
            U -= sum(self.XI[i] * np.exp(self.Q[i] * dist))
            U += sum(self.A[i] * np.exp(self.P[i] * dist))
        return U
    
    gradient = egrad(potential)

你可以通过最小化来使用它：

gp = GuptaPotential()
sol = minimize(gp.potential, x0, method='BFGS', jac=gp.gradient, options={'gtol':1e-8, 'disp':True})