rubencart:

The first and main question is whether repeating torch.autograd.functional.jacobian like in my post theoretically gives me the correct 3rd order derivatives?


I think this should give the correct results, althought you could check this with doing the Jacobian call twice and checking with torch.autograd.functional.Hessian (although Hessian doesn’t have batch-support, so you’d have to apply it iteratively)



 rubencart:

The 2nd question, the underflow problem, I could perhaps solve by giving lambda x: torch.exp(scalar_func(x)) as argument instead of scalar_func, since scalar_func returns the log of the partition function? Or does this not make sense, for numerical stability?


Applying the exp of a log value can aid the stability but it generally does become function specific. For example, in my work I have the ratio of two quantities and I can decompose that into the exp of a substraction within the log-domain.