def add_reorganize_data(
        self, name, input_name, output_name, mode="SPACE_TO_DEPTH", block_size=2
    ):
        """
        Add a data reorganization layer of type "SPACE_TO_DEPTH" or "DEPTH_TO_SPACE".
        Refer to the specification (NeuralNetwork.proto) for more details.

        Parameters
        ----------
        name: str
            The name of this layer.

        input_name: str
            The input blob name of this layer.
        output_name: str
            The output blob name of this layer.

        mode: str

            - If mode == 'SPACE_TO_DEPTH': data is moved from the spatial to the channel dimension.
              Input is spatially divided into non-overlapping blocks of size block_size X block_size
              and data from each block is moved to the channel dimension.
              Output CHW dimensions are: [C * block_size * block_size, H/block_size, C/block_size].

            - If mode == 'DEPTH_TO_SPACE': data is moved from the channel to the spatial dimension.
              Reverse of the operation 'SPACE_TO_DEPTH'.
              Output CHW dimensions are: [C/(block_size * block_size), H * block_size, C * block_size].

            - If mode == 'PIXEL_SHUFFLE':  data is moved from the channel to the spatial dimension.
              Reverse of the operation 'SPACE_TO_DEPTH'.
              Output CHW dimensions are: [C/(block_size * block_size), H * block_size, C * block_size].

        block_size: int
            Must be greater than 1. Must divide H and W, when mode is 'SPACE_TO_DEPTH'. (block_size * block_size)
            must divide C when mode is 'DEPTH_TO_SPACE' or 'PIXEL_SHUFFLE'.

        See Also
        --------
        add_flatten, add_reshape
        """
        spec_layer = self._add_generic_layer(name, [input_name], [output_name])
        spec_layer_params = spec_layer.reorganizeData

        # Set the parameters
        if block_size < 2:
            raise ValueError(
                "Invalid block_size value %d. Must be greater than 1." % block_size
            )
        spec_layer_params.blockSize = block_size

        mode = mode.upper() if isinstance(mode, str) else mode
        if mode == "SPACE_TO_DEPTH":
            spec_layer_params.mode = (
                _proto.NeuralNetwork_pb2.ReorganizeDataLayerParams.ReorganizationType.Value(
                    "SPACE_TO_DEPTH"
                )
            )
        elif mode == "DEPTH_TO_SPACE":
            spec_layer_params.mode = (
                _proto.NeuralNetwork_pb2.ReorganizeDataLayerParams.ReorganizationType.Value(
                    "DEPTH_TO_SPACE"
                )
            )
        elif mode == "PIXEL_SHUFFLE":
            if self.spec and (
                not self.spec.specificationVersion
                or self.spec.specificationVersion < _SPECIFICATION_VERSION_IOS_14
            ):
                self.spec.specificationVersion = _SPECIFICATION_VERSION_IOS_14
            spec_layer_params.mode = (
                _proto.NeuralNetwork_pb2.ReorganizeDataLayerParams.ReorganizationType.Value(
                    "PIXEL_SHUFFLE"
                )
            )
        else:
            raise NotImplementedError("Unknown reorganization mode %s." % mode)
        return spec_layer

        Add a data reorganization layer of type "SPACE_TO_DEPTH" or "DEPTH_TO_SPACE".
        Refer to the specification (NeuralNetwork.proto) for more details.

        Parameters
        ----------
        name: str
            The name of this layer.

        input_name: str
            The input blob name of this layer.
        output_name: str
            The output blob name of this layer.

        mode: str

            - If mode == 'SPACE_TO_DEPTH': data is moved from the spatial to the channel dimension.
              Input is spatially divided into non-overlapping blocks of size block_size X block_size
              and data from each block is moved to the channel dimension.
              Output CHW dimensions are: [C * block_size * block_size, H/block_size, C/block_size].

            - If mode == 'DEPTH_TO_SPACE': data is moved from the channel to the spatial dimension.
              Reverse of the operation 'SPACE_TO_DEPTH'.
              Output CHW dimensions are: [C/(block_size * block_size), H * block_size, C * block_size].

            - If mode == 'PIXEL_SHUFFLE':  data is moved from the channel to the spatial dimension.
              Reverse of the operation 'SPACE_TO_DEPTH'.
              Output CHW dimensions are: [C/(block_size * block_size), H * block_size, C * block_size].

        block_size: int
            Must be greater than 1. Must divide H and W, when mode is 'SPACE_TO_DEPTH'. (block_size * block_size)
            must divide C when mode is 'DEPTH_TO_SPACE' or 'PIXEL_SHUFFLE'.

        See Also
        --------
        add_flatten, add_reshape
        
coremltools/models/neural_network/builder.py