Jax：对向量值参数的索引求导答案

【问题标题】：Jax: Take derivative with respect to index of vector-valued argumentJax：对向量值参数的索引求导
【发布时间】：2021-11-22 15:26:48
【问题描述】：

Jax 是否支持使用 w.r.t 的派生词？向量值变量的索引？考虑这个例子（a 是一个向量/数组）：

def test_func(a):
  return a[0]**a[1]

我可以将参数编号传递给grad(..)，但我似乎无法像上面的示例那样传递向量值参数的索引。我尝试传递一个元组的元组，即，

grad(test_func, argnums=((0,),))

但这不起作用。

【问题讨论】：

标签： python jax autodiff

【解决方案1】：

没有内置的变换可以对数组的某些元素进行渐变，但是您可以通过将数组拆分为单个元素的包装函数直接执行此操作；例如：

import jax
import jax.numpy as jnp

def test_func(a):
  return a[0]**a[1]

a = jnp.array([1.0, 2.0])
fgrad = jax.grad(lambda *args: test_func(jnp.array(args)), argnums=0)
print(fgrad(*a))
# 2.0

如果您想对所有个输入分别进行渐变（返回每个条目的渐变向量），您可以使用jax.jacobian：

print(jax.jacobian(test_func)(a))
# [2. 0.]

【讨论】：

非常好，这直接解决了我的问题。我没有想过将它包装在一个 lambda 中，该 lambda 接受一个列表并用一个数组调用函数。正如您所指出的，当您需要差分 w.r.t 时，使用 jax.jacobian 非常棒。所有参数（我将来可能需要）。