Python JAX/Autograd 向量值函数的雅可比行列式
Jacobian determinant of vector-valued function with Python JAX/Autograd
我有一个将向量映射到向量的函数
我想计算它的Jacobian determinant
,
雅可比矩阵定义为
。
因为我可以使用 numpy.linalg.det
, to compute the determinant, I just need the Jacobian matrix. I know about numdifftools.Jacobian
, but this uses numerical differentiation and I'm after automatic differentiation. Enter Autograd
/JAX
(我暂时坚持使用 Autograd
,它具有 autograd.jacobian()
方法,但我很乐意使用 JAX
只要我得到我想要的)。 如何通过向量值函数正确使用此 autograd.jacobian()
函数?
举个简单的例子,我们来看函数
![f(x)=(x_0^2, x_1^2)](https://chart.googleapis.com/chart?cht=tx&chl=f(x%29%20%3D%20(x_0%5E2%2C%20x_1%5E2%29 )
它有雅可比
![J_f = diag(2 x_0, 2 x_1)](https://chart.googleapis.com/chart?cht=tx&chl=J_f%20%3D%20%5Coperatorname%7Bdiag%7D(2x_0%2C%202x_1%29 )
导致雅可比行列式
>>> 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()
return这样的嵌套数组?如果我正确地重塑它,我会得到正确的结果:
>>> 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)
这里显然两个形状都是错误的,正确的(如我想要的雅可比矩阵)会
[[[ 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]
"How do I use this autograd.jacobian()-function correctly with a vector-valued function?"
你写了
x = np.array([[3],[11]])
这有两个问题。首先是这是一个向量的向量,而 autograd 是为向量到向量的函数而设计的。第二个是 autograd 需要浮点数,而不是整数。如果您尝试根据 int 进行微分,则会出现错误。您没有看到向量向量的错误,因为 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)
"How do I need to use autograd.jacobian (or jax.jacfwd/jax.jacrev) in order to make it handle multiple vector inputs correctly?"
有一种方法可以原地做,叫做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)
我有一个将向量映射到向量的函数
我想计算它的Jacobian determinant
雅可比矩阵定义为
因为我可以使用 numpy.linalg.det
, to compute the determinant, I just need the Jacobian matrix. I know about numdifftools.Jacobian
, but this uses numerical differentiation and I'm after automatic differentiation. Enter Autograd
/JAX
(我暂时坚持使用 Autograd
,它具有 autograd.jacobian()
方法,但我很乐意使用 JAX
只要我得到我想要的)。 如何通过向量值函数正确使用此 autograd.jacobian()
函数?
举个简单的例子,我们来看函数
![f(x)=(x_0^2, x_1^2)](https://chart.googleapis.com/chart?cht=tx&chl=f(x%29%20%3D%20(x_0%5E2%2C%20x_1%5E2%29 )
它有雅可比
![J_f = diag(2 x_0, 2 x_1)](https://chart.googleapis.com/chart?cht=tx&chl=J_f%20%3D%20%5Coperatorname%7Bdiag%7D(2x_0%2C%202x_1%29 )
导致雅可比行列式
>>> 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()
return这样的嵌套数组?如果我正确地重塑它,我会得到正确的结果:
>>> 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)
这里显然两个形状都是错误的,正确的(如我想要的雅可比矩阵)会
[[[ 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]
"How do I use this autograd.jacobian()-function correctly with a vector-valued function?"
你写了
x = np.array([[3],[11]])
这有两个问题。首先是这是一个向量的向量,而 autograd 是为向量到向量的函数而设计的。第二个是 autograd 需要浮点数,而不是整数。如果您尝试根据 int 进行微分,则会出现错误。您没有看到向量向量的错误,因为 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)
"How do I need to use autograd.jacobian (or jax.jacfwd/jax.jacrev) in order to make it handle multiple vector inputs correctly?"
有一种方法可以原地做,叫做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)