python - 使用 Python JAX/Autograd 的向量值函数的雅可比行列式

Question

我有一个将向量映射到向量的函数

我想计算它的雅可比行列式

,

其中雅可比定义为

.

因为我可以使用numpy.linalg.det, 来计算行列式，所以我只需要雅可比矩阵。我知道numdifftools.Jacobian，但这使用数值微分，我在自动微分之后。输入Autograd/ JAX（我暂时坚持Autograd，它有一个autograd.jacobian()方法，但我很乐意使用JAX，只要我得到我想要的）。如何autograd.jacobian()正确使用这个函数和向量值函数？

作为一个简单的例子，我们来看看函数


具有雅可比行列式


导致雅可比行列式

>>> import autograd.numpy as np
>>> import autograd as ag
>>> x = np.array([[3],[11]])
>>> result = 4*x[0]*x[1]
array([132])
>>> jac = ag.jacobian(f)(x)
array([[[[ 6],
         [ 0]]],


       [[[ 0],
         [22]]]])
>>> jac.shape
(2, 1, 2, 1)
>>> np.linalg.det(jac)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.8/site-packages/autograd/tracer.py", line 48, in f_wrapped
    return f_raw(*args, **kwargs)
  File "<__array_function__ internals>", line 5, in det
  File "/usr/lib/python3.8/site-packages/numpy/linalg/linalg.py", line 2113, in det
    _assert_stacked_square(a)
  File "/usr/lib/python3.8/site-packages/numpy/linalg/linalg.py", line 213, in _assert_stacked_square
    raise LinAlgError('Last 2 dimensions of the array must be square')
numpy.linalg.LinAlgError: Last 2 dimensions of the array must be square

第一种方法给了我正确的值，但形状错误。为什么会.jacobian()返回这样一个嵌套数组？如果我正确地重塑它，我会得到正确的结果：

>>> jac = ag.jacobian(f)(x).reshape(-1,2,2)
array([[[ 6,  0],
        [ 0, 22]]])
>>> np.linalg.det(jac)
array([132.])

但是现在让我们看看这在数组广播中是如何工作的，当我尝试评估多个值的雅​​可比行列式时x

>>> x = np.array([[3,5,7],[11,13,17]])
array([[ 3,  5,  7],
       [11, 13, 17]])
>>> result = 4*x[0]*x[1]
array([132, 260, 476])
>>> jac = ag.jacobian(f)(x)
array([[[[ 6,  0,  0],
         [ 0,  0,  0]],

        [[ 0, 10,  0],
         [ 0,  0,  0]],

        [[ 0,  0, 14],
         [ 0,  0,  0]]],


       [[[ 0,  0,  0],
         [22,  0,  0]],

        [[ 0,  0,  0],
         [ 0, 26,  0]],

        [[ 0,  0,  0],
         [ 0,  0, 34]]]])
>>> jac = ag.jacobian(f)(x).reshape(-1,2,2)
>>> jac
array([[[ 6,  0],
        [ 0,  0]],

       [[ 0,  0],
        [ 0, 10]],

       [[ 0,  0],
        [ 0,  0]],

       [[ 0,  0],
        [14,  0]],

       [[ 0,  0],
        [ 0,  0]],

       [[ 0, 22],
        [ 0,  0]],

       [[ 0,  0],
        [ 0,  0]],

       [[26,  0],
        [ 0,  0]],

       [[ 0,  0],
        [ 0, 34]]])
>>> jac.shape
(9,2,2)

显然，这两种形状都是错误的，正确的（如我想要的雅可比矩阵）woule be

[[[ 6,  0],
  [ 0, 22]],
 [[10,  0],
  [ 0, 26]],
 [[14,  0],
  [ 0, 34]]]

和shape=(6,2,2)

我需要如何使用autograd.jacobian（或jax.jacfwd/ jax.jacrev）才能使其正确处理多个向量输入？

---

注意：使用显式循环并手动处理每个点，我得到了正确的结果。但是有没有办法做到这一点？

>>> dets = []
>>> for v in zip(*x):
>>>    v = np.array(v)
>>>    jac = ag.jacobian(f)(v)
>>>    print(jac, jac.shape, '\n')
>>>    det = np.linalg.det(jac)
>>>    dets.append(det)
 [[ 6.  0.]
 [ 0. 22.]] (2, 2)

 [[10.  0.]
 [ 0. 26.]] (2, 2)

 [[14.  0.]
 [ 0. 34.]] (2, 2)

>>> dets
 [131.99999999999997, 260.00000000000017, 475.9999999999998]

Answer 1

“如何正确使用这个 autograd.jacobian() 函数和向量值函数？”

你写过

x = np.array([[3],[11]])

这有两个问题。首先是这是一个向量的向量，而 autograd 是为向量到向量函数而设计的。第二个是 autograd 需要浮点数，而不是整数。如果你试图区分整数，你会得到一个错误。您没有看到向量向量的错误，因为 autograd 会自动将您的整数列表转换为浮点数列表。

TypeError: Can't differentiate w.r.t. type <class 'int'>

下面的代码应该给你决定因素。

import autograd.numpy as np
import autograd as ag

def f(x):
    return np.array([x[0]**2,x[1]**2])

x = np.array([3.,11.])
jac = ag.jacobian(f)(x)
result = np.linalg.det(jac)
print(result)

“我如何需要使用 autograd.jacobian（或 jax.jacfwd/jax.jacrev）才能正确处理多个向量输入？”

有一种方法可以做到这一点，它被称为 jax.vmap。请参阅 JAX 文档。（https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap）

在这种情况下，我可以使用以下代码计算雅可比行列式向量。请注意，我可以用与以前完全相同的方式定义函数 f，vmap 在幕后为我们完成工作。

import jax.numpy as np
import jax

def f(x):
    return np.array([x[0]**2,x[1]**2])

x = np.array([[3.,11.],[5.,13.],[7.,17.]])

jac = jax.jacobian(f)
vmap_jac = jax.vmap(jac)
result = np.linalg.det(vmap_jac(x))
print(result)