Return only the diagonal of A @ B as a 1-D tensor, without materializing the full n×n product matrix.
A @ B
n×n
Sign in to write and run your own code.
solve(...)