我试图了解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
?另外,我可以使用相同的函数来获取 Hessian 矩阵吗?
我查看了有关它的JAX 帮助部分,但无法弄清楚