Return the diagonal of the matrix product A @ B without computing the full (n, n) product first.
A @ B
(n, n)
solve(...)