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.
        Refer to the ``FlattenLayerParams`` message in the
        specification (NeuralNetwork.proto) for more details.

        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_layer = self._add_generic_layer(name, [input_name], [output_name])
        spec_layer_params = spec_layer.flatten

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

        return spec_layer

        Add a flatten layer. Only flattens the channel, height and width axis.
        Leaves the sequence axis as is.
        Refer to the ``FlattenLayerParams`` message in the
        specification (NeuralNetwork.proto) for more details.

        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
        
coremltools/models/neural_network/builder.py