class_docstr: Temporal Entropy Model
class_docstr_tok: ["Temporal","Entropy","Model"]
docstr: Args:
      latent_unquantized: the latent to transmit (quantize), expected shape is
          [B, C, H, W]
      previous_latents: previously transmitted (quantized) latents, should be of
          size at least one and at most `context_len`. Each PreviousLatent has
              - quantized: floats (i.e. noised) tensor of shape [B, C, H, W]
              - processed: [b', seq_len_enc, d_model]
  
  Returns:
      TemporalEntropyModelOut, see docstring there.
docstr_tok: ["Args",":","latent_unquantized",":","the","latent","to","transmit","(","quantize",")",",","expected","shape","is","[","B",",","C",",","H",",","W","]","previous_latents",":","previously","transmitted","(","quantized",")","latents",",","should","be","of","size","at","least","one","and","at","most","`","context_len","`",".","Each","PreviousLatent","has","-","quantized",":","floats","(","i.e",".","noised",")","tensor","of","shape","[","B",",","C",",","H",",","W","]","-","processed",":","[","b","'",",","seq_len_enc",",","d_model","]","Returns",":","TemporalEntropyModelOut",",","see","docstring","there","."]
code: def forward(
      self,
      latent_unquantized: Tensor,
      previous_latents: Sequence[PreviousLatent],
  ) -> TemporalEntropyModelOut:
      """
      Args:
          latent_unquantized: the latent to transmit (quantize), expected shape is
              [B, C, H, W]
          previous_latents: previously transmitted (quantized) latents, should be of
              size at least one and at most `context_len`. Each PreviousLatent has
                  - quantized: floats (i.e. noised) tensor of shape [B, C, H, W]
                  - processed: [b', seq_len_enc, d_model]
  
      Returns:
          TemporalEntropyModelOut, see docstring there.
      """
      H, W = latent_unquantized.shape[-2:]
      # encoded_seqs: list of tensors [b', seq_len_enc=patch_size_enc^2, d_model]
      encoded_seqs = self._get_encoded_seqs(
          previous_latents=previous_latents
      )
      b_enc, _, d_enc = encoded_seqs[
          0
      ].shape  # [b', patch_size_enc^2, d_enc]
      if d_enc != self.d_model:
          raise ValueError(f"Shape mismatch, {d_enc}!={self.d_model}")
  
      # Soft rounding via straight-through gradient estimation
      latent_q = self.round_st(
          latent_unquantized
      )  # [B, C, H, W], float32 of ints
  
      # Patcher expects [B, C, H, W] and returns [b', seq_len, C]
      latent_q_patched, (n_h, n_w) = self.patcher(
          latent_q, self.window_size_dec
      )
      (
          b_dec,
          seq_len,
          d_dec,
      ) = latent_q_patched.shape  # b_dec, patch_size_dec^2, C
      if d_dec != self.num_channels:
          raise ValueError(
              f"Model dims don't match, {d_dec}!={self.num_channels}"
          )
      if b_dec != b_enc:
          raise ValueError(
              f"Batch dims don't match, got {b_enc} != {b_dec}!"
          )
      assert seq_len == self.window_size_dec**2, "Error patching"
  
      # Transformer expects inputs to have channels in the last dim, [b', seq_len, C]
      # Transformer expects quantized latents
      mean, scale, dec_output = self._get_transformer_output(
          encoded_patched=torch.cat(
              encoded_seqs, dim=-2
          ),  # cat on seq dim
          latent_q_patched=latent_q_patched,  # [b_dec, seq_len, C], ROUNDED
      )
      assert (
          mean.shape == latent_q_patched.shape
      ), f"Shape mismatch! {mean.shape} != {latent_q_patched.shape}"
  
      decoder_features = self.patcher.unpatch(
          x_patched=(dec_output, (n_h, n_w)),
          crop=(H, W),
          channels_last=True,
      )  # [B, d_model, H, W]
  
      latent_unquantized_patched = self.patcher(
          latent_unquantized, self.window_size_dec
      ).tensor  # [b', seq_len_dec, C]
  
      assert (
          latent_unquantized_patched.shape == scale.shape
      )  # [B', seq_len_dec, C]
      # We use GaussianConditional from compressai as bottleneck, its forward method
      # returns (output, likelihood)
      # in VCT scales and means are quantized in the bottleneck
      output, likelihood = self.bottleneck(
          inputs=latent_unquantized_patched, scales=scale, means=mean
      )  # output is a noised version of the `inputs`
      assert output.shape == likelihood.shape  # [b', seq_len, C]
  
      output = self.patcher.unpatch(
          (output, (n_h, n_w)), crop=(H, W), channels_last=True
      )
      assert output.shape == latent_unquantized.shape
  
      return TemporalEntropyModelOut(
          perturbed_latent=output,
          bits=likelihood,
          features=decoder_features,
      )
  
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 7 for calling the VCTEntropyModel obj's underlying member method code with example usage: obj.forward(latent_unquantized, previous_latents) and returns: TemporalEntropyModelOut