human
2. SoftMax

input:   batch_size x n_feats
output: batch_size x n_feats

$\text{softmax}(x)_i = \frac{\exp x_i} {\sum_j \exp x_j}$
Recall that $\text{softmax}(x) == \text{softmax}(x - \text{const})$. It makes possible to avoid computing exp() from large argument."
Please write code following the instructions in jupyter notebook style.
assistant
#Here is code to 2.

class SoftMax(Module):
    def __init__(self):
         super(SoftMax, self).__init__()
    
    def updateOutput(self, input):
        # start with normalization for numerical stability
        self.output = np.subtract(input, input.max(axis=1, keepdims=True))
        # Your code goes here. ################################################
        self.output = np.exp(self.output)/np.sum(np.exp(self.output), axis=1, keepdims=True)
        return self.output
    
    def updateGradInput(self, input, gradOutput):
        # Your code goes here. ################################################
        pred_grad = np.multiply(self.output, gradOutput)
        self.gradInput = pred_grad - np.multiply(self.output, np.sum(pred_grad, axis=1, keepdims=True))
        return self.gradInput
    
    def __repr__(self):
        return "SoftMax"