<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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.")