Objective

PyTorch currently does not have a way to take diagonals with respect to arbitrary axes as e.g. numpy.diagonal has.
Two immediate use cases are

• Taking the diagonal of a batch of matrices. This is handy e.g. in Gaussian processes, but I have seen this a number of times.
• One of the missing einsum features (torch.einsum('jii->ji', a)) would benefit from this as well.

Plan

I propose to do the following (see https://github.com/t-vi/pytorch/tree/diagonal_with_dim for a preliminary implementation)

• Implement diagonal natively in ATen instead of referring to TH/THC diag.
• Make diagonal return a view by adjusting shape, stride, and offset. diag copies, but this seems impractical for higher dimensional arrays and also unneeded. This is not backward compatible.
• Add two dimension parameters dim1 and dim2.
• Make a diagonal_backward method.
• Include the Tensor method, too.

Edit: Based on your feedback, I updated the code to use numpy semantics.

Limiting numpy compatibility?

NumPy's implementation has two features that might be worth to differ from:

• NumPy's default axes are 0 and 1. I think it is more natural to use -2 and -1 ("batch thinking").
• NumPy moves the new axis replacing the two old ones at the end of the tensor. I think it might be more natural to put it into the place of the first axis to be removed. I must admit that this is more a gut feeling than based on hard facts.

I would greatly appreciate your input on this, in particular regarding the potential numpy deviations.