Evaluating a section of a jacobian

Let's say that I have some function f(x), where x is a vector and returns a vector. I can evaluate the Jacobian of this function fairly simply, as demonstrated below.

from autograd import numpy as np
import autograd as ag
def f(x):
    return np.array([x.sum(),(x[:3]**2).sum(),np.log(np.exp(x).sum())])
xtest=np.array([0,.5,.3,.2])
print(f(xtest))

print(ag.jacobian(f)(xtest))

My question is if there's some way of evaluating only some columns of this jacobian. For example, let's say I only wanted the first and last columns of it. So far I haven't found any way of evaluating this more efficiently than just evaluating the whole jacobian and throwing some away. If anyone can help please let me know!

Yes, it is possible to evaluate only certain columns of the Jacobian matrix using Autograd. Autograd provides a jacobian_vector_product function that allows you to compute the product of the Jacobian matrix with a given vector. By setting the input vector to be a standard basis vector, you can obtain the corresponding column of the Jacobian matrix.

from autograd import numpy as np
import autograd as ag

def f(x):
    return np.array([x.sum(), (x[:3]**2).sum(), np.log(np.exp(x).sum())])

xtest = np.array([0, 0.5, 0.3, 0.2])
J = ag.jacobian(f)(xtest)  # Compute the full Jacobian matrix
e1 = np.zeros_like(xtest)
e1[0] = 1  # Set the first component to 1
e3 = np.zeros_like(xtest)
e3[-1] = 1  # Set the last component to 1

col1 = ag.jacobian_vector_product(f, xtest)(e1)  # Compute first column
col3 = ag.jacobian_vector_product(f, xtest)(e3)  # Compute third column

print(col1)
print(col3)

we compute the Jacobian matrix of f with respect to xtest using ag.jacobian(f)(xtest). We then create two vectors e1 and e3, which are the standard basis vectors corresponding to the first and last columns of the Jacobian matrix, respectively. We compute these columns by calling ag.jacobian_vector_product(f, xtest)(e1) and ag.jacobian_vector_product(f, xtest)(e3), respectively.

HIPS / autograd

Evaluating a section of a jacobian #582