Pass X through a softmax function in a numerically stable way using the
    log-sum-exp trick.

    This transformation is:

    .. math::

        \\frac{\exp\{X_k\}}{\sum^K_{j=1} \exp\{X_j\}}

    and is appliedx to each row/column, `k`, of X.

    Parameters
    ----------
        X: ndarray
            2D array of shape (N, D) to apply the log-sum-exp trick.
        axis: int, optional
            Axis to apply the summation along (works the same as axis in
            numpy.sum).

    Returns
    -------
        smX: ndarray
            results of applying the log-sum-exp trick, this will be shape
            (N, D), and each row will sum to 1 if :code:`axis=1` or each column
            will sum to 1 if :code:`axis=0`.

def softmax(X, axis=0):
    """
    Pass X through a softmax function in a numerically stable way using the
    log-sum-exp trick.

    This transformation is:

    .. math::

        \\frac{\exp\{X_k\}}{\sum^K_{j=1} \exp\{X_j\}}

    and is appliedx to each row/column, `k`, of X.

    Parameters
    ----------
        X: ndarray
            2D array of shape (N, D) to apply the log-sum-exp trick.
        axis: int, optional
            Axis to apply the summation along (works the same as axis in
            numpy.sum).

    Returns
    -------
        smX: ndarray
            results of applying the log-sum-exp trick, this will be shape
            (N, D), and each row will sum to 1 if :code:`axis=1` or each column
            will sum to 1 if :code:`axis=0`.
    """

    if axis == 1:
        return np.exp(X - logsumexp(X, axis=1)[:, np.newaxis])
    elif axis == 0:
        return np.exp(X - logsumexp(X, axis=0))
    else:
        raise ValueError("This only works on 2D arrays for now.")