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