import torch


def softmax(y, dim=0, beta=30.0):
    """
    differentiable substitute for argmax(y)
    
    max(y) = softmax(y)*y
    
    Example
    -------
    >>> y = torch.tensor([0.7, 0.99, 0.9, 0.1])
    >>> softmax(y)
    tensor([1.5607e-04, 9.3688e-01, 6.2963e-02, 2.3770e-12])
    >>> y * softmax(y)
    tensor([1.0925e-04, 9.2751e-01, 5.6667e-02, 2.3770e-13])
    """
    a = torch.exp(beta*y)
    b = torch.sum(a, dim).unsqueeze(dim)
    return a/b