Return only the diagonal of A @ B as a 1-D tensor, without materializing the full n×n product matrix.

Input

A =
tensor([[0., 1., 2.],
        [3., 4., 5.],
        [6., 7., 8.]])
B =
tensor([[ 9., 10., 11.],
        [12., 13., 14.],
        [15., 16., 17.]])

Output

tensor([ 42., 162., 300.])