Row-wise Jacobian with pytorch

Question

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[0]
     d_i = X.shape[1]

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

     y = f(x)
     d_o = y.shape[1]

     v = ?
     y.backward(v)

     Z = x.grad
     return Z

Answer 1

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.