Have problem when trying calculate jacobian matrix from nonlinear system of equations

Question

I found code for compute jacobian matrix from here and try it for non-linear system of equations.

import autograd.numpy as np
from autograd import jacobian

x = np.array([1,2], dtype=float)

def fs(x,y):
  return np.array([x + 2*y - 2, x**2 + 4*y**2 - 4])

jacobian_cost = jacobian(fs)

jacobian_cost(x[0],x[1])

But instead of outputting result like this

array([ 1.  2.]
 [ 2. 16.])

it outputting result like this

array([1., 2.])

I curious about whats wrong with this code, maybe i missing something.

Answer 1

You need to vectorize your function, and call it in a vectorized manner. So define

def fs(x):
    return np.stack([x[0,:]+ 2*x[1, :]-2, x[0, :]**2+4*x[1,:]**2-4])

and then call

x = np.array([[1.0, 2]])
output = jacobian(fs)(x)

But note that the autograd library is very outdated, doesn't have a very good documentation and is no longer maintained. I'd recommend using PyTorch or JAX instead.