# Row-wise Jacobian with pytorch

Suppose I have $$f:\mathbb{R}^{d_i}\to\mathbb{R}^{d_o}$$. Let $$X \in \mathbb{R}^{n \times d_i}$$ and I apply $$f$$ to each row of $$X$$, obtaining $$Y = f(X) \in \mathbb{R}^{n \times d_o}$$. I would like to compute a tensor $$Z$$ which is defined as

$$Z_{i,j,k} = \frac{\partial Y_{i,j}}{\partial X_{i,k}}$$

using pytorch. How should I define v to achieve that?

def get_row_wise_jacobian(X, f):
n = X.shape
d_i = X.shape

x = X.clone().detach().requires_grad_(True)

y = f(x)
d_o = y.shape

v = ?
y.backward(v)

Z = x.grad
return Z


## 1 Answer

If f(x) is a deep network, you could use the differential network implementation of this repository. This differential network adds a few computations to the feed-forward graph to compute the analytical Jacobian w.r.t. to the network input. In contrast to the autograd implementation, the computational overhead is minimal and the Jacobian is computed using forward differentiation. Furthermore, the examples of the Repository show that you want to use softplus activation to get smoother Jacobians.