def _reshape(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())

        Reshapes the matrix W as expained in step (1).
        
understand_icl/fairseq/fairseq/modules/quantization/pq/pq.py