class_docstr: Temporal Entropy Model
class_docstr_tok: ["Temporal","Entropy","Model"]
docstr: Temporal Entropy Model
  Args:
      num_channels: number of channels in the latent space,
          i.e. symbols per token. Defaults to 196.
      context_len: number of previous latents. Defaults to 2.
      window_size_enc: window (patch) size in encoder.
          Defaults to 8.
      window_size_dec: window (patch) size in decoder.
          Defaults to 4.
      num_layers_encoder_sep: number of layers in the separate encoder.
          Defaults to 3.
      num_layers_encoder_joint: number of layers in the joint encoder.
          Defaults to 2.
      num_layers_decoder: number of layers in the decoder.
          Defaults to 5.
      d_model: feature dimensionality inside the model.
          Defaults to 768.
      num_head: number of attention heads in MHA layers.
          Defaults to 16.
      mlp_expansion: expansion *factor* for each MLP.
          Defaults to 4.
      drop_out_enc: dropout probability in encoder.
          Defaults to 0.0.
      drop_out_dec: dropout probability in decoder.
          Defaults to 0.0.
docstr_tok: ["Temporal","Entropy","Model","Args",":","num_channels",":","number","of","channels","in","the","latent","space",",","i.e",".","symbols","per","token",".","Defaults","to","196.","context_len",":","number","of","previous","latents",".","Defaults","to","2.","window_size_enc",":","window","(","patch",")","size","in","encoder",".","Defaults","to","8.","window_size_dec",":","window","(","patch",")","size","in","decoder",".","Defaults","to","4.","num_layers_encoder_sep",":","number","of","layers","in","the","separate","encoder",".","Defaults","to","3.","num_layers_encoder_joint",":","number","of","layers","in","the","joint","encoder",".","Defaults","to","2.","num_layers_decoder",":","number","of","layers","in","the","decoder",".","Defaults","to","5.","d_model",":","feature","dimensionality","inside","the","model",".","Defaults","to","768.","num_head",":","number","of","attention","heads","in","MHA","layers",".","Defaults","to","16.","mlp_expansion",":","expansion","*","factor","*","for","each","MLP",".","Defaults","to","4.","drop_out_enc",":","dropout","probability","in","encoder",".","Defaults","to","0.0.","drop_out_dec",":","dropout","probability","in","decoder",".","Defaults","to","0.0","."]
code: def __init__(
      self,
      num_channels: int = 192,
      context_len: int = 2,
      window_size_enc: int = 8,
      window_size_dec: int = 4,
      num_layers_encoder_sep: int = 3,
      num_layers_encoder_joint: int = 2,
      num_layers_decoder: int = 5,
      d_model: int = 768,
      num_head: int = 16,
      mlp_expansion: int = 4,
      drop_out_enc: float = 0.0,
      drop_out_dec: float = 0.0,
  ) -> None:
      """
      Temporal Entropy Model
      Args:
          num_channels: number of channels in the latent space,
              i.e. symbols per token. Defaults to 196.
          context_len: number of previous latents. Defaults to 2.
          window_size_enc: window (patch) size in encoder.
              Defaults to 8.
          window_size_dec: window (patch) size in decoder.
              Defaults to 4.
          num_layers_encoder_sep: number of layers in the separate encoder.
              Defaults to 3.
          num_layers_encoder_joint: number of layers in the joint encoder.
              Defaults to 2.
          num_layers_decoder: number of layers in the decoder.
              Defaults to 5.
          d_model: feature dimensionality inside the model.
              Defaults to 768.
          num_head: number of attention heads in MHA layers.
              Defaults to 16.
          mlp_expansion: expansion *factor* for each MLP.
              Defaults to 4.
          drop_out_enc: dropout probability in encoder.
              Defaults to 0.0.
          drop_out_dec: dropout probability in decoder.
              Defaults to 0.0.
      """
      super().__init__()
      if window_size_enc < window_size_dec:
          raise ValueError(
              f"window_size_enc={window_size_enc} cannot be lower"
              f"than window_size_dec={window_size_dec}."
          )
      if num_channels < 0:
          raise ValueError(
              f"num_channels={num_channels} cannot be negative"
          )
      self.num_channels = num_channels
      self.window_size_enc = window_size_enc
      self.window_size_dec = window_size_dec
  
      self.d_model = d_model
      # we will use compressai's GsnConditional as a bottleneck
      self.bottleneck = GsnConditionalLocScaleShift(
          num_scales=256,
          num_means=100,
          min_scale=0.01,
          tail_mass=(2 ** (-8)),
      )
  
      self.range_bottleneck = None
  
      self.context_len = context_len
      self.encoder_sep = EncoderSection(
          num_layers=num_layers_encoder_sep,
          num_heads=num_head,
          hidden_dim=d_model,
          mlp_expansion=mlp_expansion,
          dropout=drop_out_enc,
      )
      self.encoder_joint = EncoderSection(
          num_layers=num_layers_encoder_joint,
          num_heads=num_head,
          hidden_dim=d_model,
          mlp_expansion=mlp_expansion,
          dropout=drop_out_enc,
      )
  
      self.seq_len_dec = window_size_dec**2
      self.seq_len_enc = window_size_enc**2
  
      self.patcher = Patcher(window_size_dec, "reflect")
      self.learned_zero = StartSym(hidden_dim=num_channels)
  
      self.enc_position_sep = LearnedPosition(
          seq_length=self.seq_len_enc, hidden_dim=d_model
      )
      self.enc_position_joint = LearnedPosition(
          seq_length=self.seq_len_enc * context_len, hidden_dim=d_model
      )
  
      self.dec_position = LearnedPosition(
          seq_length=self.seq_len_dec, hidden_dim=d_model
      )
  
      self.post_embedding_layernorm = nn.LayerNorm(d_model, eps=1e-6)
  
      self.encoder_embedding = make_embedding(
          input_dim=num_channels, hidden_dim=d_model
      )  # a single linear layer
      self.decoder_embedding = make_embedding(
          input_dim=num_channels, hidden_dim=d_model
      )  # a single linear layer
  
      self.decoder = Transformer(
          seq_length=self.seq_len_dec,
          num_layers=num_layers_decoder,
          num_heads=num_head,
          hidden_dim=d_model,
          mlp_expansion=mlp_expansion,
          dropout=drop_out_dec,
          is_decoder=True,
      )  # its forward method returns [B', seq_length, hidden_dim]
  
      def _make_final_heads(output_channels: int) -> nn.Module:
          # 3 stacked linear layers with leakyrelu activations
          return nn.Sequential(
              nn.Linear(d_model, d_model),
              nn.LeakyReLU(),
              nn.Linear(d_model, d_model),
              nn.LeakyReLU(),
              nn.Linear(d_model, output_channels),
          )
  
      self.mean_head = _make_final_heads(num_channels)
      self.scale_head = _make_final_heads(num_channels)
  
imports: import itertools
  import sys
  from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple
  import torch
  import torch.nn
  from torch import Tensor
  from .bottlenecks import GsnConditionalLocScaleShift
  from .entropy_model_layers import EncoderSection, LearnedPosition, StartSym, Transformer
  from .layers_utils import make_embedding
  from .patcher import Patcher
label_desc: Use image node_id 1 to create a new VCTEntropyModel object from inherited base classes: nn with example: obj = VCTEntropyModel(num_channels, context_len, window_size_enc, window_size_dec, num_layers_encoder_sep, num_layers_encoder_joint, num_layers_decoder, d_model, num_head, mlp_expansion, drop_out_enc, drop_out_dec)