Let's suppose I have the code below and I want to calculate the jacobian of L, which is the prediction made by a neural network in Pytorch, L is of size nx1 where n is the number of samples in a mini batch. In order to avoid a for loop for each entry of L (n entries) to calculate the jacobian for each sample in the mini batch some codes I found just sum the n predictions of the neural network (L) with respect with the inputs and then calculate the gradient of the sum. First I can't understand why is the gradient of the sum the same of the sum of the gradients for each sample in pytorch architecture. Second I tried both with the sum and with a for loop and the results diverge. Could it be due to numerical approximations or because the sum just doesn't make sense?
The code is below, where both functions belong to a nn.module:
def forward(self, x):
with torch.set_grad_enabled(True):
def function(x,t):
self.n = n = x.shape[1]//2
qqd = x.requires_grad_(True)
L = self._lagrangian(qqd).sum()
J = grad(L, qqd, create_graph=True)[0]
def _lagrangian(self, qqd):
x = F.softplus(self.fc1(qqd))
x = F.softplus(self.fc2(x))
x = F.softplus(self.fc3(x))
L = self.fc_last(x)
return L