def METHOD_NAME(self, W):
    """
    Reshapes the matrix W as expained in step (1).
    """
    # fully connected: by convention the weight has size out_features x in_features
    if len(W.size()) == 2:
        self.out_features, self.in_features = W.size()
        assert (
            self.in_features % self.block_size == 0
        ), "Linear: n_blocks must be a multiple of in_features"
        return (
            W.reshape(self.out_features, -1, self.block_size)
            .permute(2, 1, 0)
            .flatten(1, 2)
        )
    # convolutional: we reshape along the spatial dimension
    elif len(W.size()) == 4:
        self.out_channels, self.in_channels, self.k_h, self.k_w = W.size()
        assert (
            self.in_channels * self.k_h * self.k_w
        ) % self.block_size == 0, (
            "Conv2d: n_blocks must be a multiple of in_channels * k_h * k_w"
        )
        return (
            W.reshape(self.out_channels, -1, self.block_size)
            .permute(2, 1, 0)
            .flatten(1, 2)
        )
    # not implemented
    else:
        raise NotImplementedError(W.size())