Differentiable top-k function Is there any differentiable function that, for a given vector, selects and encourages the top-k maximum value and suppresses the rest of the values? For example for z = [0.01 0.1 0.04 0.5 0.24] the top 3 would something like this:
top-3(z) = [1e-10 0.89 2e-9 0.98 0.92] 

 A: Below is a simple differentiable top-k for PyTorch I wrote.
See also the code here.
The idea is to pick some sigmoid function, e.g. the simple $\sigma(x) = \frac{1}{1 + e^{-x}}$, and express your probabilities by $z = \sigma(x)$.
In your case that would mean
$$x = [-4.5951, -2.1972, -3.1781, -0.0000, -1.1527].$$
Now we pick a value $t\in\mathbb R$ such that $\sum_i \sigma(x_i + t) = k$.
In your case we would choose roughly $t=2.8301$, which would result in the "modified distribution":
$$\sigma(x+t) = [0.1462, 0.6531, 0.4139, 0.9443, 0.8426],$$
which you can check sums to 3.
Notes
There is a lot of flexibility in which $\sigma(x)$ function you choose, which can give you more or less "aggressive encouragement".
For example, taking $\sigma(x) = \frac{1}{\sqrt{1+e^{-x}}}$ gives the output
$[0.09, 0.67, 0.34, 0.98, 0.91]$.
Another nice function is $\sigma(x) = \exp(\min(x, 0))$, which has the nice property of coinciding with the standard softmax function at $k=1$. In your case it gives $[0.07, 0.67, 0.27, 1.00, 1.00]$. Note however that it is not completely differentiable.
Math
The main question remains: Is this function differentiable? And if so, is it efficiently differentiable?
Let $t(x)$ be the $t$ value needed to sum to $k$ for a given set of $x$ values.
Let $f(x) \sigma(x + t(x))$ and consider
$$
\frac{df(x)_i}{dx_j} 
= \frac{d\sigma(x_i+t(x))}{dx_j}
= \sigma'(x_i + t)\left([i=j] + \frac{dt(x)}{d x_j}\right).
$$
Issue: How do we determine $\frac{dt(x)}{d x_j}$?
Turns out this is easy using Implicit differentiation.
We simply take the derivative of the equation $k=\sum_i \sigma(x_i+t)$,
which gives us the equation
$$
\frac{dk}{d x_j}
= 0
= \sum_i \sigma'(x_i + t(x))\left([i=j] + \frac{dt(x)}{d x_j}\right)
= \sigma'(x_j+t(x)) + \frac{dt(x)}{d x_j}\sum_i \sigma'(x_i + t(x)).
$$
Solving we find
$
\frac{dt(x)}{d x_j} = \frac{-\sigma'(x_j+t(x))}{\sum_i \sigma'(x_i + t(x))}.
$
That's all there is to it.
If we define $v=\sigma'(x+t_k(x))$, the Jacobian of $\text{top}_k$ at $x$ is $J_{\text{top}_k}(x) = \text{diag}(v) - vv^T/\|v\|_1$.
Given $t$ this is easy to compute and the vector-Jacobian product $u^T J_{\text{top}_k}(x)$, needed for backpropagation, is just $u\circ v - \langle u, v\rangle v/\|v\|_1$.
This is what's implemented below and verified numerically.
To actually compute $t(x)$, I use a binary search, taking advantage of the $\sigma$ functions being monotone.
A nice property of the problem is that we only have to compute $t$ on the forward pass. When computing the gradient we simple reuse the $t$ we already computed.
Code
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):
        ts, ps = _find_ts(xs, k)
        ctx.save_for_backward(xs, ts)
        return ps

    @staticmethod
    def backward(ctx, grad_output):
        # Compute vjp, that is grad_output.T @ J.
        xs, ts = ctx.saved_tensors
        # Let v = sigmoid'(x + t)
        v = sigmoid_grad(xs + ts)
        s = v.sum(dim=1, keepdims=True)
        # Jacobian is -vv.T/s + diag(v)
        uv = grad_output * v
        t1 = - uv.sum(dim=1, keepdims=True) * v / s
        return t1 + uv, None

@torch.no_grad()
def _find_ts(xs, k):
    b, n = xs.shape
    assert 0 < k < n
    # Lo should be small enough that all sigmoids are in the 0 area.
    # Similarly Hi is large enough that all are in their 1 area.
    lo = -xs.max(dim=1, keepdims=True).values - 10
    hi = -xs.min(dim=1, keepdims=True).values + 10
    for _ in range(64):
        mid = (hi + lo)/2
        mask = sigmoid(xs + mid).sum(dim=1) < k
        lo[mask] = mid[mask]
        hi[~mask] = mid[~mask]
    ts = (lo + hi)/2
    return ts, sigmoid(xs + ts)

topk = TopK.apply
xs = torch.randn(2, 3)
ps = topk(xs, 2)
print(xs, ps, ps.sum(dim=1))

from torch.autograd import gradcheck
input = torch.randn(20, 10, dtype=torch.double, requires_grad=True)
for k in range(1, 10):
    print(k, gradcheck(topk, (input, k), eps=1e-6, atol=1e-4))

