Compatibility problem with pytorch trace

[jplumail]

torch.linalg.trace should work on stacks of matrices.

But this

import array_api_compat.torch as torch

a = torch.reshape(torch.arange(10*5*5), (10,5,5))
torch.linalg.trace(a)

throws RuntimeError: trace: expected a matrix, but got tensor with dim 3.

The pytorch's trace function doesn't handle stacks: see the docs.

[asmeurer]

Thanks for the report. What's a good way to work around this? @lezcano

[lezcano]

x.diagonal(0, -2, -1).sum(-1)

[asmeurer]

This was fixed in #80