"autograd.grad", one per Jacobian row. Using "vmap()", we can
   vectorize the whole computation, computing the Jacobian in a single
   call to "autograd.grad".



Setup
N = 5
f = lambda x: x ** 2
x = torch.randn(N, requires_grad=True)
y = f(x)
I_N = torch.eye(N)
Sequential approach
jacobian_rows = [torch.autograd.grad(y, x, v, retain_graph=True)[0]
                 for v in I_N.unbind()]
jacobian = torch.stack(jacobian_rows)
vectorized gradient computation
def get_vjp(v):
    return torch.autograd.grad(y, x, v)
jacobian = torch.vmap(get_vjp)(I_N)



"vmap()" can also be nested, producing an output with multiple
   batched dimensions



torch.dot                            # [D], [D] -> []
batched_dot = torch.vmap(torch.vmap(torch.dot))  # [N1, N0, D], [N1, N0, D] -> [N1, N0]
x, y = torch.randn(2, 3, 5), torch.randn(2, 3, 5)
batched_dot(x, y) # tensor of size [2, 3]


