How to sum over specifi samples in batch

1

I have the following batch

tensor([[[ 0.,  1.],
         [ 2.,  3.],
         [ 4.,  5.],
         [ 6.,  7.]],

        [[ 8.,  9.],
         [10., 11.],
         [12., 13.],
         [14., 15.]],

        [[16., 17.],
         [18., 19.],
         [20., 21.],
         [22., 23.]]])

with shape [3,4, 2]. how can I get different summation combinations along the batch dimension? for instance I want the first row of output to be summation of row1 and row3 and second row, just row 2, i.e.

tensor([[[ 16.,  18.],
         [ 20.,  22.],
         [ 24.,  26.],
         [ 28.,  30.]],

        [[ 8.,  9.],
         [10., 11.],
         [12., 13.],
         [14., 15.]],

I thought of multiplication with tensor tensor([[1., 0., 1.],
[0., 1., 0.]]), but obviously there is dimension mismatch. Thank you.

2

I figured out the solution. torch.matmul(m.T, n.T).T, where m is the main tensor and n is the tensor of 0 and 1s. Just in case someone needs it.