import torch
import math


def softmax(scores: torch.Tensor, base: float = math.e, axis: int = -1) -> torch.Tensor:
    """Returns softmax array for array of scores

    Converts a set of raw scores from a model (logits) into a
    probability distribution via softmax.

    The probability distribution will be a set of real numbers
    such that each is in the range 0-1.0 and the sum is 1.0.

    Assumes input is a pytorch tensor: tensor([1.0, 4.0, 2.0, 3.0])

    Args
        prediction (pytorch tensor) a pytorch tensor of any positive/negative real numbers.
        base (float) the base for the exponential (default e)
    """
    exps = base ** scores.float()  # exponential for each value in array
    sum_exps = torch.sum(exps, dim=axis, keepdim=True)  # sum of all exponentials
    prob_dist = exps / sum_exps  # normalize exponentials

    return prob_dist