Определитель Якоби вектор-функции с Python JAX / Autograd

Question

У меня есть функция, которая отображает векторы на векторы

R^n">

, и я хочу вычислить ее определитель Якоби

,

, где якобиан определяется как

.

Поскольку я могу использовать numpy.linalg.det, чтобы вычислить определитель, мне просто нужна матрица Якоби. Я знаю о numdifftools.Jacobian, но для этого используется численное дифференцирование, и я после автоматического c дифференцирования. Введите 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() возвращает такой вложенный массив? Если я изменяю форму это правильно, я получаю правильный результат:

>>> 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) для того, чтобы заставить его правильно обрабатывать несколько векторных входов?

---

Примечание. Используя явный l oop и обрабатывая каждую точку вручную, я получаю правильный результат. Но есть ли способ сделать это на месте?

>>> 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)