Get the jacobian of a vector-valued function that takes batch inputs, in pytorch.

torch_jacobian.py

def get_jacobian(net, x, noutputs):
    x = x.squeeze()
    n = x.size()[0]
    x = x.repeat(noutputs, 1)
    x.requires_grad_(True)
    y = net(x)
    y.backward(torch.eye(noutputs))
    return x.grad.data

import torch
from torch.autograd import grad
from torch import autograd

def get_jacobian(net, x, noutputs):
    x = x.unsqueeze(1) # b, 1 ,in
    #print(x.shape)
    n = x.size()[0]
    x = x.repeat(1, noutputs, 1) # out,in => b, out, in
    #print(x.shape)
    x.requires_grad_(True)
    y = net(x)
    #print(y.shape) # b, o, o
    input_val = torch.eye(noutputs).reshape(1,noutputs, noutputs).repeat(n, 1, 1)
    #print(input_val.shape)
    y.backward(input_val)
    #print(y.shape)
    return x.grad.data


class Net(torch.nn.Module):
    def __init__(self, in_dim, out_dim):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(in_dim, out_dim)
    
    def forward(self, x):
        return self.fc1(x)    

batch = 2
num_features = 3
num_outputs = 5
x = torch.ones(batch, num_features)
print(x.shape)
print(net(x).shape)
net = Net(num_features, num_outputs)

get_jacobian(net, x, num_outputs).shape