Row-wise Jacobian with pytorch

Suppose I have $f:mathbb{R}^{d_i}tomathbb{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