Autograd中避免数组赋值：含数组赋值的目标函数求导问题

我明白你遇到的问题——Autograd确实没法追踪普通的数组赋值操作，因为这类操作会打断它依赖的计算图链路。针对你的obj函数，我们可以通过以下几个关键调整来实现对theta的正确求导：

1. 用向量化操作替代显式赋值/循环 #

你的代码里先初始化了res = np.zeros(n)，后续大概率是通过循环给res的每个元素赋值。这种逐元素赋值是Autograd无法追踪的，我们需要把这部分逻辑改成向量化的numpy运算，直接生成结果数组，跳过初始化再赋值的步骤。

举个例子，如果你的Score计算是每个res[i]等于theta[i]的平方和，原来的代码可能是：

for i in range(n):
    res[i] = np.sum(theta[i] ** 2)

现在可以直接改成：

res = np.sum(theta ** 2, axis=1)

完全不需要初始化res再赋值，Autograd能完美追踪这种向量化操作。

2. 替换原生numpy为Autograd的numpy实现 #

Autograd提供了自己的numpy封装（autograd.numpy），它重写了numpy的核心操作来支持求导追踪。把函数里的np替换成autograd.numpy，确保所有数组操作都被Autograd监控。

调整后的基础函数框架会是这样：

import autograd.numpy as ag_np

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 = ag_np.atleast_2d(theta)
    # 直接用向量化运算生成res，替换你的Score逻辑
    res = ag_np.sum(theta ** 2, axis=1)  # 示例操作，换成你的实际计算
    return res

3. 处理条件赋值场景（如果有的话） #

如果你的Score逻辑里有条件判断式的赋值，比如“当theta[i,0] > 0时res[i]取某个值，否则取另一个值”，不要用循环+if-else赋值，而是用ag_np.where实现向量化的条件选择：

# 示例：根据theta的第一列取值生成res
res = ag_np.where(theta[:, 0] > 0, theta[:, 0], ag_np.zeros_like(theta[:, 0]))

核心原则 #

Autograd只能追踪返回新数组的纯函数式操作，任何原地修改数组（比如res[i] = x）、或者破坏计算图的操作都会导致求导失败。所以只要把所有赋值/循环逻辑转化为向量化的numpy运算，就能让Autograd正确计算theta的导数。

内容的提问来源于stack exchange，提问作者p-value