了解JAX argnums在其梯度函数中的参数

Question

我试图了解argnums在JAX的梯度函数中的行为。假设我有以下函数:

def make_mse(x, t):  
def mse(w,b): 
return np.sum(jnp.power(x.dot(w) + b - t, 2))/2
return mse 

我用下面的方法进行渐变:

w_gradient, b_gradient = grad(make_mse(train_data, y), (0,1))(w,b)

argnums= (0,1)在这种情况下，但它是什么意思?梯度是根据哪些变量计算的?如果我用argnums=0代替会有什么不同?同样，我可以用同样的函数得到黑森矩阵吗?

我看了一下JAX帮助部分，但是不明白

Answer 1

当您向grad传递多个argnums时，结果是一个返回梯度元组的函数，相当于您分别计算每个梯度:

def f(x, y):
return x ** 2 + x * y + y ** 2
df_dxy = grad(f, argnums=(0, 1))
df_dx = grad(f, argnums=0)
df_dy = grad(f, argnums=1)
x = 3.0
y = 4.25
assert df_dxy(x, y) == (df_dx(x, y), df_dy(x, y))

如果你想计算混合二阶导数，你可以通过重复应用梯度来做到这一点:

d2f_dxdy = grad(grad(f, argnums=0), argnums=1)
assert d2f_dxdy(x, y) == 1