ML Katas

by: newest upvotes saves

Extract Diagonal from a Batch of Matrices

this year hard (<10 mins) | linear algebra einsum

### Description Given a batch of square matrices of shape `(B, N, N)`, extract the diagonal of each matrix. The output should be a tensor of shape `(B, N)`. This can be achieved with...

1