Calculating Jacobian using functorch

Hello, I am trying to calculate the Jacobian matrix of a GAN using functorch, so I need the derivative of each intermediate layer with respect to the input (z), I am using the generator in this tutorial, for example, I need to calculate the derivative of the 5th layer of the generator with respect to z, to do this I am using the chain rule in the following loop:

grads = []
z = sample_batch(20) #Generates a random vector of size (20,100)
inp = z[:,:,None,None] #Reshaping to match the Convolution layer
index = 5
grad = 1

for i in range(index):
  jac_func = lambda x: G.main[i](x)
  if i != 0 :
    inp = Get_Hidden_Output(z, i-1)  #Gets the output of the i-1 layer when the input of the model is z
  grad = vmap(jac_func)(inp)*grad

However, while computing, I get an error when reaching to the BatchNorm layers of the model, so I would like to know if it is possible to calculate the such operation directly without using the chain rule, any help would be appreciated, just in case the error I get for the BatchNorm layer is:

ValueError: expected 4D input (got 3D input)
1 Like

Do you get the same error when you use a 4D input?

No, I used z.unsqueeze(1) to avoid the error, but I am not convinced with the solution.