了解其渐变函数中的 JAX argnums 参数答案

【问题标题】：Understanding JAX argnums parameter in its gradient function了解其渐变函数中的 JAX argnums 参数
【发布时间】：2021-03-02 19:18:06
【问题描述】：

我试图了解 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 help 部分的相关内容，但无法弄清楚

【问题讨论】：

标签： python automatic-differentiation jax

【解决方案1】：

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

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

【讨论】：