test.py

import torch
import torch.nn.functional as F


def topk_softmax(logits, k):
    vals, idx = torch.topk(logits, k, dim=-1)
    return idx, F.softmax(vals, dim=-1)


def full_softmax_then_renorm(logits, k):
    probs = F.softmax(logits, dim=-1)
    idx = torch.topk(logits, k, dim=-1).indices
    p = probs.gather(dim=-1, index=idx)
    w = p / p.sum(dim=-1, keepdim=True)
    return idx, w


torch.manual_seed(0)

T, E, k = 4, 8, 2
logits1 = torch.randn(T, E, dtype=torch.double, requires_grad=True)
logits2 = logits1.detach().clone().requires_grad_()

idx1, w1 = topk_softmax(logits1, k)
idx2, w2 = full_softmax_then_renorm(logits2, k)

assert torch.equal(idx1, idx2)
print("forward max diff:", (w1 - w2).abs().max().item())

# Some arbitrary downstream gradient.
grad_out = torch.randn_like(w1)

loss1 = (w1 * grad_out).sum()
loss2 = (w2 * grad_out).sum()

loss1.backward()
loss2.backward()

print("backward max diff:", (logits1.grad - logits2.grad).abs().max().item())

"""
forward max diff: 1.1102230246251565e-16
backward max diff: 1.1102230246251565e-16
"""