class_docstr: Temporal Entropy Model
class_docstr_tok: ["Temporal","Entropy","Model"]
docstr: Compress and decompress autoregressively. Can only handle batch size 1.
  
  Args:
      latent_unquantized: unquantized latent of shape [1, C, H, W]
      previous_latents: a sequence of length at least 1 and at most `context_len`
          containingg PreviousLatent objects which hold 2 tensors:
              - quantized: [1, C, H, W]
              - processed: latents passed through the encoder, expected shape
                  [b', seq_len_enc, d_model]
      run_decode: bool, defaults to False. Whether to run the actual decoding
  
  Returns:
      TemporalEntropyModelOut object with the following components:
          - perturbed_latent: tensor of ints with same shape as the input tensor
              `latent_unquatnized` -- [1, C, H, W]
          - bits: number of bits used to compress the input latent, float tensor
          - features: features tensor from the decoder, shape [1, d_model, H, W]
docstr_tok: ["Compress","and","decompress","autoregressively",".","Can","only","handle","batch","size","1",".","Args",":","latent_unquantized",":","unquantized","latent","of","shape","[","1",",","C",",","H",",","W","]","previous_latents",":","a","sequence","of","length","at","least","1","and","at","most","`","context_len","`","containingg","PreviousLatent","objects","which","hold","2","tensors",":","-","quantized",":","[","1",",","C",",","H",",","W","]","-","processed",":","latents","passed","through","the","encoder",",","expected","shape","[","b","'",",","seq_len_enc",",","d_model","]","run_decode",":","bool",",","defaults","to","False",".","Whether","to","run","the","actual","decoding","Returns",":","TemporalEntropyModelOut","object","with","the","following","components",":","-","perturbed_latent",":","tensor","of","ints","with","same","shape","as","the","input","tensor","`","latent_unquatnized","`","--","[","1",",","C",",","H",",","W","]","-","bits",":","number","of","bits","used","to","compress","the","input","latent",",","float","tensor","-","features",":","features","tensor","from","the","decoder",",","shape","[","1",",","d_model",",","H",",","W","]"]
code: def compress(
      self,
      *,
      latent_unquantized: Tensor,
      previous_latents: Sequence[PreviousLatent],
      run_decode: bool = False,
      validate_causal: bool = False,
  ) -> TemporalEntropyModelOut:
      """
      Compress and decompress autoregressively. Can only handle batch size 1.
  
      Args:
          latent_unquantized: unquantized latent of shape [1, C, H, W]
          previous_latents: a sequence of length at least 1 and at most `context_len`
              containingg PreviousLatent objects which hold 2 tensors:
                  - quantized: [1, C, H, W]
                  - processed: latents passed through the encoder, expected shape
                      [b', seq_len_enc, d_model]
          run_decode: bool, defaults to False. Whether to run the actual decoding
  
      Returns:
          TemporalEntropyModelOut object with the following components:
              - perturbed_latent: tensor of ints with same shape as the input tensor
                  `latent_unquatnized` -- [1, C, H, W]
              - bits: number of bits used to compress the input latent, float tensor
              - features: features tensor from the decoder, shape [1, d_model, H, W]
      """
      B, C, H, W = latent_unquantized.shape
      assert B == 1, "Cannot handle batch yet."
  
      encoded_seqs = self._get_encoded_seqs(previous_latents)
      # previously coded latents, shape is [b', 2*seq_len_enc, d_model], floats
      encoded = torch.cat(encoded_seqs, -2)
  
      latent_patched, (n_h, n_w) = self.patcher(
          latent_unquantized, self.window_size_dec
      )  # [b', seq_len, C]
  
      if validate_causal:
          self.validate_causal(
              latent_q_patched=latent_patched, encoded=encoded
          )
  
      # Encoding: compress to strings - strings is a list of len seq_len_dec (16)
      strings, extra = self._encode(latent_patched, encoded)
      means, scales, dec_output, quantized = extra.values()
  
      decoder_features = self.patcher.unpatch(
          (dec_output, (n_h, n_w)), crop=(H, W), channels_last=True
      )  # [1, d_model, H, W]
  
      # Count bits in each sequence, each string is a list of len 1
      bits = [sum(len(string[0]) * 8 for string in strings)]
  
      # decoding:
      if not run_decode:
          # For performance, since coding is lossless, real decode can be skipped
          decoded = torch.round(latent_unquantized)
      else:
          use_output_from_encode = True
          decoded = self._decode(
              strings,
              encoded,
              shape=(H, W, C),
              encode_means=means,
              encode_scales=scales,
              use_output_from_encode=use_output_from_encode,
          )
          dequantized = self.patcher.unpatch(
              x_patched=(quantized, (n_h, n_w)),
              crop=(H, W),
              channels_last=True,
          )
          if use_output_from_encode:
              # This should pass if `use_output_from_encode=True`
              assert (
                  decoded == dequantized
              ).all(), "Something went wrong!"
  
      return TemporalEntropyModelOut(
          perturbed_latent=decoded,
          bits=torch.tensor(bits, dtype=torch.float32),
          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 10 for calling the VCTEntropyModel obj's underlying member method code with example usage: obj.compress() and returns: TemporalEntropyModelOut