<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def add_flatten(self, name, mode, input_name, output_name):
 """ Add a flatten layer. Only flattens the channel, height and width axis. Leaves the sequence axis as is. Parameters name: str The name of this layer. mode: int - If mode == 0, the flatten layer is in CHANNEL_FIRST mode. - If mode == 1, the flatten layer is in CHANNEL_LAST mode. input_name: str The input blob name of this layer. output_name: str The output blob name of this layer. See Also -------- add_permute, add_reshape """

        spec = self.spec
        nn_spec = self.nn_spec

        # Add a new layer
        spec_layer = nn_spec.layers.add()
        spec_layer.name = name
        spec_layer.input.append(input_name)
        spec_layer.output.append(output_name)
        spec_layer_params = spec_layer.flatten

        # Set the parameters
        if mode == 0:
            spec_layer_params.mode = \
                        _NeuralNetwork_pb2.FlattenLayerParams.FlattenOrder.Value('CHANNEL_FIRST')
        elif mode == 1:
            spec_layer_params.mode = \
                        _NeuralNetwork_pb2.FlattenLayerParams.FlattenOrder.Value('CHANNEL_LAST')
        else:
            raise NotImplementedError(
                'Unknown flatten mode %d ' % mode)