I am trying to use autograd to calculate the product of a Jacobian matrix and a vector, but could not make it work efficiently. Any suggestion will be highly appreciated:
x is a
M dimensional vector. I have N functions
y1, y2, ... yn such that
y1 = y1(x), y2 = y2(x), ..., yn = yn(x). The derivative of
y with respect to
x then form a
N x M Jacobian matrix. Let’s call the Jacobian matrix
I am interested in calculating the product between
v, a M-dimensional vector, i.e.
Jv. Since autograd gives us the opportunity to calculate gradients, I thought it can be done. What I realize that is that, while gradients were computed, the backward() immediately adds up
dyn/dx1 before I can have a chance to calculate Jv.
The only way I can see is to build a loop to iterate
y1, y2, ..., yn, which is quite inefficient (in my opinion). I am looking for more efficient ways. Any feedback is welcome.
P.S. the python code below shows a simple example with M = 3 and N = 2 linear functions. Since the function is linear, the Jacobian is essentially the weight matrix. The desirable output is thus product of weight and v.
# Inputs dim_input = 3 # number of scalar functions num_function = 2 x = nd.arange(dim_input) print("input x:", x) # linear weights for every function weights = nd.arange(num_function*dim_input).reshape(shape=(num_function, -1)) - 3.5 print("weights:", weights) v = nd.arange(dim_input)+100 print("vector v:", v) x.attach_grad() with autograd.record(): y = nd.dot(x, weights.T) y.backward() grad = x.grad print("function output: ", y) print("Gradient -- cannot get Jacobian: ", grad) print("Target result:", nd.dot(v.T, weights.T))