How to select top half of the elements in each row of a given tensor? Pytorch code
I have a tensor A with shape (Batch, sequence, dimension).
I calculated the importance weights for each element in a sequence. The weights can be represented as a tensor B with shape (Batch, sequence).
I need to select the first half of the elements with large weights. That means the result is expected to have a shape (Batch, sequence/2, dimension).
Note:

Each sequence of tensor A has different length. I pad the sequences to form the tensor A.
I have the mask matrix Mask_A with shape (Batch, sequence).
Therefore, I can't simply take the top sequence/2 elements because of padding.
I know it can be achieved by 'for' loop in python code. However, it will slow down the program. I hope it can be done in Pytorch functions to get GPU acceleration.
The result doesn't have to be sorted according to weights.
It's also acceptable to receive a result of mask containing top-half elements.

Here is an example code:
A = torch.randn((2, 4, 3))
# batch size = 2, sequence legnth = 4, dimension = 3 
# it looks like 
'''tensor([[[ 6.0840e-01,  4.5604e-01, -1.3264e+00],
         [-4.6437e-01,  1.6999e-01,  1.3551e+00],
         [-1.9888e+00, -2.3047e-01,  1.2347e-03],
         [ 0.0000e+00,  0.0000e+00,  0.0000e+00]],

        [[ 7.9035e-01, -5.5752e-01, -1.2477e+00],
         [-1.7801e-01,  4.6232e-01,  1.3019e+00],
         [ 0.0000e+00,  0.0000e+00,  0.0000e+00],
         [ 0.0000e+00,  0.0000e+00,  0.0000e+00]]])'''
# [ 0.0000e+00,  0.0000e+00,  0.0000e+00] means padding

Mask_A = torch.tensor([[1., 1., 1., 0.], [1., 1., 0., 0.]])
# this means 1-st row of A has three elements, 2-nd row has two elements. '0.' means padding

B = torch.tensor([[0.3,0.1,0.6,0.0],[0.05, 0.95, 0.0, 0.0]])
# this is the unsorted importance weights matrix

# what I want is like:
'''
tensor([[[-1.9888e+00, -2.3047e-01,  1.2347e-03],
         [ 0.0000e+00,  0.0000e+00,  0.0000e+00]],

        [[-1.7801e-01,  4.6232e-01,  1.3019e+00],
         [ 0.0000e+00,  0.0000e+00,  0.0000e+00]]])
'''

# [ 0.0000e+00,  0.0000e+00,  0.0000e+00] means padding

# what I want can also be like:
Mask_result = torch.tensor([[0., 0., 1., 0.], [0., 1., 0., 0.]])


Making some assumptions here about exactly how many elements you want to keep, but this should work:
# Get a tensor marking the rank of each weight within the sequence
rank = torch.argsort(torch.argsort(B, axis=1, descending=True), axis=1)

# Determine how many elements per sequence we want to keep, based on the mask
n_to_keep = torch.floor(torch.sum(Mask_A, axis=1)/2)

# Make a mask marking the n elements per sequence with the top weights
# using broadcasting
mask_result = rank &lt; n_to_keep[:,None]

