examples/lm/layers.py
def compute_tensor(self, x):
        """
        :param x: (batch, time, vec)
        """
        # Target class
        class_matrix = self.target_tensor // self.output_size
        class_vector = class_matrix.reshape((-1,))
        # Target index
        target_matrix = self.target_tensor % self.output_size
        target_vector = target_matrix.reshape((-1,))
        # Input matrix
        input_matrix = x.reshape((-1, self.input_dim))
        # Output matrix
        output_tensor3d = self.output_layer.compute_tensor(x)
        output_matrix = output_tensor3d.reshape((-1, self.class_size, self.output_size))
        arange_vec = self.arange_cache[:output_matrix.shape[0]]
        sub_output_matrix = output_matrix[arange_vec, class_vector]
        # Softmax
        softmax_output_matrix = self.softmax_layer.compute_tensor(sub_output_matrix)
        # Class prediction
        class_output_matrix = self.class_layer.compute_tensor(x)
        # Costs
        output_cost = LMCost(softmax_output_matrix, target_vector).get()
        class_cost = LMCost(class_output_matrix, class_matrix).get()
        final_cost = output_cost + class_cost

        return final_cost
def compute_tensor(self, x):
        """
        :param x: (batch, time, vec)
        """
        # Target class
        class_matrix = self.target_tensor // self.output_size
        class_vector = class_matrix.reshape((-1,))
        # Target index
        target_matrix = self.target_tensor % self.output_size
        target_vector = target_matrix.reshape((-1,))
        # Input matrix
        input_matrix = x.reshape((-1, self.input_dim))
        # Output matrix
        output_tensor3d = self.output_layer.compute_tensor(x)
        output_matrix = output_tensor3d.reshape((-1, self.class_size, self.output_size))
        arange_vec = self.arange_cache[:output_matrix.shape[0]]
        sub_output_matrix = output_matrix[arange_vec, class_vector]
        # Softmax
        softmax_output_matrix = self.softmax_layer.compute_tensor(sub_output_matrix)
        # Class prediction
        class_output_matrix = self.class_layer.compute_tensor(x)
        # Costs
        output_cost = LMCost(softmax_output_matrix, target_vector).get()
        class_cost = LMCost(class_output_matrix, class_matrix).get()
        final_cost = output_cost + class_cost

        return final_cost
:param x: (batch, time, vec)