How to calculate a jacobian for an entire batch

[KFrank]

Hi Yedidya!

@Yedidya_kfir

I believe the jacobian calculate the gradient of every batch w.r.t. to every input (ignoring the separation between batches)

How can I get the function to treat those batches as separate without using a for loop and separating them myself?

If I understand what you want to do, I believe that @AlphaBetaGamma96’s
post that uses vmap() and jacrev() will give you your “batch jacobian:”

Computing batch Jacobian efficiently

Just an update for anyone who reads the thread in the future, as of PyTorch2,the functorch library is now included in pytorch. So you can replace functorch with torch.func, for the most part the syntax is the same except if you have an nn.Module you’ll need to create a ‘functional’ version of your model. For example, model = myModel(*args, **kwargs) #our network from torch.func import vmap, jacrev, functional_call params = dict(model.named_parameters()) inputs = torch.randn(batch_size, input…

Best.

K. Frank