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

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

v = ?
y.backward(v)