class SoftmaxCrossent(Module):

    def run(self, x, labels):
        # Calculate exp(x-max(x))
        x -= x.max(1, keepdims=True)
        sum_e_x = np.exp(x).sum(1, keepdims=True)

        # Calculate log(softmax(x))
        log_soft = x - np.log(sum_e_x)
        crossed = - np.multiply(labels, log_soft)

        # Cache values
        self.cache_value(log_soft=log_soft, labels=labels)
        return crossed.sum(1).mean()

    def backprop(self, grad=1.0):
        # Get cached values
        cache = self.get_cached_values()
        labels, log_soft = cache['labels'], cache['log_soft']
        batch_size = labels.shape[0]

        # backprop through mean() == divide by batch_size
        scaling = grad / batch_size
        return scaling * (np.exp(log_soft) - labels)