How should I compute the Jacobian matrix of a layer with respect its input?

hello guys
I’m trying to compute the Jacobian of a hidden layer with respect to its input so I found `torch.autograd.functional.vjp’, but I’m not sure that is what I am searching about that.
‘vjp’ gives three input arguments (func, input,v) the first argument in my problem is the function which produces the input of mentioned hidden layer (e.g. f(x*w)+b), the second argument is “x” and the last is the hidden layer’s weight matrix.
shall I use ‘vjp’? any help would be appreciated.