human
Write me a code to 3 given the following:
3. LogSoftMax

input:   batch_size x n_feats
output: batch_size x n_feats

$\text{logsoftmax}(x)_i = \log\text{softmax}(x)_i = x_i - \log {\sum_j \exp x_j}$
The main goal of this layer is to be used in computation of log-likelihood loss."
Please write code following the instructions in jupyter notebook style.
assistant
#Here is code to 3.

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