martinsipka:

def batch_jacobian(f, x):
    f_sum = lambda x: torch.sum(f(x), axis=0)
    return jacobian(f_sum, x)



Very nice, didn’t think of it. Yields the exact same results as using a for loop, but runs much faster than the for loop.