I’m trying to implement efficient k-winner-take-all activation function in PyTorch.
class KWinnersTakeAll(torch.autograd.Function):

    @staticmethod
    def forward(ctx, tensor, sparsity: float):
        batch_size, embedding_size = tensor.shape
        _, argsort = tensor.sort(dim=1, descending=True)
        k_active = math.ceil(sparsity * embedding_size)
        active_indices = argsort[:, :k_active]
        mask_active = torch.ByteTensor(tensor.shape).zero_()
        for sample_id in range(batch_size):
            mask_active[sample_id, active_indices[sample_id]] = 1
        tensor[~mask_active] = 0
        tensor[mask_active] = 1
        ctx.save_for_backward(mask_active)
        return tensor

What I want is to get rid of for-loop by vectorizing the computations it does. Any idea how to do that?
Let mask_active.shape == (256, 128) and active_indices.shape == (256, 7) then
mask_active[active_indices].shape == (256, 7, 128) but I’d like to have mask_active[active_indices].shape == (256, 7) so then I could assign mask_active[active_indices] = 1