thomasahle

import torch
from functorch import vmap, grad
from torch.autograd import Function

sigmoid = torch.sigmoid
sigmoid_grad = vmap(vmap(grad(sigmoid)))

class TopK(Function):
    @staticmethod
    def forward(ctx, xs, k):