Log-sum-exp trick for matrix X for summation along a specified axis.

    This performs the following operation in a stable fashion,

    .. math::

        \log \sum^K_{k=1} \exp\{x_k\}

    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
    -------
        lseX: ndarray
            results of applying the log-sum-exp trick, this will be shape (D,)
            if :code:`axis=0` or shape (N,) if :code:`axis=1`.

def logsumexp(X, axis=0):
    """
    Log-sum-exp trick for matrix X for summation along a specified axis.

    This performs the following operation in a stable fashion,

    .. math::

        \log \sum^K_{k=1} \exp\{x_k\}

    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
    -------
        lseX: ndarray
            results of applying the log-sum-exp trick, this will be shape (D,)
            if :code:`axis=0` or shape (N,) if :code:`axis=1`.
    """

    mx = X.max(axis=axis)
    if (X.ndim > 1):
        mx = np.atleast_2d(mx).T if axis == 1 else np.atleast_2d(mx)

    return np.log(np.exp(X - mx).sum(axis=axis)) + np.ravel(mx)