Python 避免autograd中的数组赋值

我从autograd教程中了解到，当阵列包含在要区分的目标中时，不支持阵列分配。然而，目前我的代码中有以下目标函数，我想就θ进行区分：

def obj(theta):
    """
    Computes the objective function to be differentiated.

    Args:
        theta: np.array of shape (n, d)

    Return:
        res: np.array of shape (n,)
    """
    theta = np.atleast_2d(theta)
    n = theta.shape[0]

    res = np.zeros(n)  # Scores
    for i in xrange(n):
        res[i] = ... # Do some computations with theta[i, :]

    return res

通常我可以通过在θ上矢量化计算来避免for循环；然而，在这种情况下，计算已经涉及到各种线性代数运算、逆运算等。给定一行θ作为超参数，我发现很难对所有θ行进行向量化运算。在这种情况下，我不知道还有比用for循环逐行填充res数组更好的方法

我尝试了一种简单的方法来避免数组赋值，方法是创建一个列表，并在每次迭代时将结果附加到该列表中，然后在返回res时最终将该列表转换为数组，但最终我得到了所有零梯度

我想知道在此设置中推荐的一般解决方案是什么？

您可以使用numpy.apply\u沿\u轴为数据中的特定轴应用函数

def func(row):
    # return the computation result for "row"

def obj(theta):
    """
    Computes the objective function to be differentiated.

    Args:
        theta: np.array of shape (n, d)

    Return:
        res: np.array of shape (n,)
    """
    theta = np.atleast_2d(theta)
    n = theta.shape[0]

    res = np.apply_along_axis(func1d=func, axis=1, arr=a)

    return res

---

谢谢你的建议，但是我得到了错误键error:in self.vjp=primitive_vjps[fun]parent_argnums，value，args，kwargs。。。看来autograd也不适合这个。。。