Problem with several parameters to torch.autograd.jacobian()

autograd
1

Hi,

I want to calculate the jacobian matrix w.r.t. the parameters of a model. This is how I am doing it at the moment:

from functorch import jacrev, vmap, make_functional, grad

def _compute_centered_jacobian(model,samples):
    func, parameters = make_functional(model)
    def func_ampl(params):
        return func(params, samples)
    def func_phase(params):
        return func(params, samples)[1]
    jac_ampl = torch.autograd.functional.jacobian(func_ampl, parameters)
    jac_phase = torch.autograd.functional.jacobian(func_phase parameters)
    return jac_ampl, jac_phase

but I am getting the error TypeError: func_ampl() takes 1 positional argument but 9 were given because there are 9 different weight matrices in the model. How can I resolve this problem?
Thanks a lot for your help in advance!