Hi everyone! For my research, I have the need to compute a Jacobian-vector product (JVP), where the Jacobian is of the outputs of a nn.Module (on a mini-batch) w.r.t. its parameters. Basically, the output of the JVP has the same size as the outputs of the original network.
I installed PyTorch 1.5 because of the new functional.jvp. The input to jvp must be a function with tensor inputs/outputs (in my case: the inputs are all the Parameters of a nn.Module). However, PT does not have a functional implementation of a Module, where the call to the module takes the parameters as arguments (like in JAX).
In this case, I would need to pass the parameters explicitly, then copy them inside the module every time the function is used, which looks inefficient.
Is there a simpler way to combine the new interface with a nn.Module? Or is there a simpler way to achieve this JVP which I am overseeing? Thanks!