def _permute_output(self, x: torch.Tensor) -> torch.Tensor:
        """
        >>> BaseRNN(BaseRNNConfig(10, 16, num_heads=2, output_shape='SBH'))._permute_output(torch.zeros((5, 2, 16))).shape
        torch.Size([5, 2, 16])
        >>> BaseRNN(BaseRNNConfig(10, 16, num_heads=2, output_shape='BNSH'))._permute_output(torch.zeros((5, 3, 16))).shape
        torch.Size([3, 2, 5, 8])
        >>> BaseRNN(BaseRNNConfig(10, 16, num_heads=2, output_shape='SBNH'))._permute_output(torch.zeros((5, 3, 16))).shape
        torch.Size([5, 3, 2, 8])
        """
        if self.config.output_shape == "SBH":
            return x
        elif self.config.output_shape == "BSH":
            return x.permute(1, 0, 2)
        elif self.config.output_shape == "BNSH":
            return x.view(
                (x.shape[0], x.shape[1], self.config.num_heads, self.config.head_dim)
            ).permute(1, 2, 0, 3)
        elif self.config.output_shape == "SBNH":
            return x.view(
                (x.shape[0], x.shape[1], self.config.num_heads, self.config.head_dim)
            )

        >>> BaseRNN(BaseRNNConfig(10, 16, num_heads=2, output_shape='SBH'))._permute_output(torch.zeros((5, 2, 16))).shape
        torch.Size([5, 2, 16])
        >>> BaseRNN(BaseRNNConfig(10, 16, num_heads=2, output_shape='BNSH'))._permute_output(torch.zeros((5, 3, 16))).shape
        torch.Size([3, 2, 5, 8])
        >>> BaseRNN(BaseRNNConfig(10, 16, num_heads=2, output_shape='SBNH'))._permute_output(torch.zeros((5, 3, 16))).shape
        torch.Size([5, 3, 2, 8])
        
xlstm/blocks/slstm/cell.py