def _permute_input(self, x: torch.Tensor) -> torch.Tensor:
        """
        # TODO Adapt this
        # >>> BaseRNN(BaseRNNConfig(10, 10, num_heads=2, input_shape='SBG'))._permute_input(torch.zeros((5, 2, 10))).shape
        # torch.Size([5, 2, 10])
        # >>> BaseRNN(BaseRNNConfig(10, 10, num_heads=2, input_shape='BSG'))._permute_input(torch.zeros((5, 2, 10))).shape
        # torch.Size([2, 5, 10])
        """
        if self.config.input_shape == "SBGNH":
            y = x.view(
                x.shape[0], x.shape[1], self.config.num_gates, self.config.num_heads, -1
            )
        elif self.config.input_shape == "BSGNH":
            y = x.view(
                x.shape[0], x.shape[1], self.config.num_gates, self.config.num_heads, -1
            ).permute(1, 0, 2, 3, 4)
        else:
            raise ValueError("Bad input_shape value")
        if self.config.internal_input_shape == "SBGNH":
            return y.view(y.shape[0], y.shape[1], -1)
        elif self.config.internal_input_shape == "SBNGH":
            return y.permute(0, 1, 3, 2, 4).reshape(y.shape[0], y.shape[1], -1)
        elif self.config.internal_input_shape == "SBNHG":
            return y.permute(0, 1, 3, 4, 2).reshape(y.shape[0], y.shape[1], -1)
        else:
            raise ValueError("Bad internal_input_shape value")

        # TODO Adapt this
        # >>> BaseRNN(BaseRNNConfig(10, 10, num_heads=2, input_shape='SBG'))._permute_input(torch.zeros((5, 2, 10))).shape
        # torch.Size([5, 2, 10])
        # >>> BaseRNN(BaseRNNConfig(10, 10, num_heads=2, input_shape='BSG'))._permute_input(torch.zeros((5, 2, 10))).shape
        # torch.Size([2, 5, 10])
        
xlstm/blocks/slstm/cell.py