content: from typing import *
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from ...modules import sparse as sp
  from .base import SparseTransformerBase
  
  
  class SLatEncoder(SparseTransformerBase):
      def __init__(
          self,
          resolution: int,
          in_channels: int,
          model_channels: int,
          latent_channels: int,
          num_blocks: int,
          num_heads: Optional[int] = None,
          num_head_channels: Optional[int] = 64,
          mlp_ratio: float = 4,
          attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
          window_size: int = 8,
          pe_mode: Literal["ape", "rope"] = "ape",
          use_fp16: bool = False,
          use_checkpoint: bool = False,
          qk_rms_norm: bool = False,
      ):
          super().__init__(
              in_channels=in_channels,
              model_channels=model_channels,
              num_blocks=num_blocks,
              num_heads=num_heads,
              num_head_channels=num_head_channels,
              mlp_ratio=mlp_ratio,
              attn_mode=attn_mode,
              window_size=window_size,
              pe_mode=pe_mode,
              use_fp16=use_fp16,
              use_checkpoint=use_checkpoint,
              qk_rms_norm=qk_rms_norm,
          )
          self.resolution = resolution
          self.out_layer = sp.SparseLinear(model_channels, 2 * latent_channels)
  
          self.initialize_weights()
          if use_fp16:
              self.convert_to_fp16()
  
      def initialize_weights(self) -> None:
          super().initialize_weights()
          # Zero-out output layers:
          nn.init.constant_(self.out_layer.weight, 0)
          nn.init.constant_(self.out_layer.bias, 0)
  
      def forward(self, x: sp.SparseTensor, sample_posterior=True, return_raw=False):
          h = super().forward(x)
          h = h.type(x.dtype)
          h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
          h = self.out_layer(h)
          
          # Sample from the posterior distribution
          mean, logvar = h.feats.chunk(2, dim=-1)
          if sample_posterior:
              std = torch.exp(0.5 * logvar)
              z = mean + std * torch.randn_like(std)
          else:
              z = mean
          z = h.replace(z)
              
          if return_raw:
              return z, mean, logvar
          else:
              return z
  