I have the exact same question as OP. I tested using torch.diagonal but because I also want to compute the second-order jacobian doing this using diagonal is extremely slow.
Here is the code I’m using
    def compute_u_x(self, x):
        self.u_x = torch.autograd.functional.jacobian(self, x, create_graph=True)
        self.u_x = torch.diagonal(self.u_x, dim1=0, dim2=2).permute(2, 0, 1)
        return self.u_x
    
    def compute_u_xx(self, x):
        self.u_xx = torch.autograd.functional.jacobian(self.compute_u_x, x)
        self.u_xx = torch.diagonal(self.u_xx, dim1=0, dim2=3).permute(3, 0, 1, 2)
        self.u_xx = torch.diagonal(self.u_xx, dim1=2, dim2=3)
        return self.u_xx

This will give me what I need, but it’s computing so many extra things it significantly slowed my training loop. I also made a version which isn’t vectorized that uses the torch.autograd.grad function which works exponentially faster than above. However, this method only works for scalar output functions, which is not what I need.