class_docstr: Lossless compression for Herarchical latent variable models.
  
  Based on the Craystack (https://github.com/j-towns/craystack)
  implementation of Bits-Back coding.
  
  Args:
      model: The hierarchical latent variable model.
      data_shape: The shape [batch-size, channels, height, width]
          of the data and each latent.
      latent_prec: Coding precision.
      prior_prec: Precision for the discretization.
      obs_prec: Precision for the observation codec.
      data_prec: Precision of the discrete data.
      initial_len: Length of the random initial bitstring.
      data_min: Minimum value of the data to be fed into the model.
      data_max: Maximum value of the data to be fed into the model.
      clip_stdd: Minimum value for the stdd of the observation codec.
      track: Whether to track the message sizes at each step.
class_docstr_tok: ["Lossless","compression","for","Herarchical","latent","variable","models",".","Based","on","the","Craystack","(","https",":","\/\/github.com\/j-towns\/craystack",")","implementation","of","Bits-Back","coding",".","Args",":","model",":","The","hierarchical","latent","variable","model",".","data_shape",":","The","shape","[","batch-size",",","channels",",","height",",","width","]","of","the","data","and","each","latent",".","latent_prec",":","Coding","precision",".","prior_prec",":","Precision","for","the","discretization",".","obs_prec",":","Precision","for","the","observation","codec",".","data_prec",":","Precision","of","the","discrete","data",".","initial_len",":","Length","of","the","random","initial","bitstring",".","data_min",":","Minimum","value","of","the","data","to","be","fed","into","the","model",".","data_max",":","Maximum","value","of","the","data","to","be","fed","into","the","model",".","clip_stdd",":","Minimum","value","for","the","stdd","of","the","observation","codec",".","track",":","Whether","to","track","the","message","sizes","at","each","step","."]
docstr: Track the message length at a given time-step.
  
  Args:
      codec: Codec that adheres to the craystack api.
      step: The step in the hierarchy.
  Returns:
      A craystack codec which tracks the pops and pushes.
docstr_tok: ["Track","the","message","length","at","a","given","time-step",".","Args",":","codec",":","Codec","that","adheres","to","the","craystack","api",".","step",":","The","step","in","the","hierarchy",".","Returns",":","A","craystack","codec","which","tracks","the","pops","and","pushes","."]
code: def _track_codec(self, codec: cs.Codec, step: int) -> cs.Codec:
      """
      Track the message length at a given time-step.
  
      Args:
          codec: Codec that adheres to the craystack api.
          step: The step in the hierarchy.
      Returns:
          A craystack codec which tracks the pops and pushes.
      """
      if not self.track:
          return codec
  
      if not 0 <= step <= self.model.steps:
          raise RuntimeError(
              f"Step {step} is out of bounds [0, {self.model.steps}]."
          )
  
      def track_push(message, data):
          init_bytes = cs.flatten(message).nbytes
          (message,) = codec.push(message, data)
          self.push_bytes[step] += (
              cs.flatten(message).nbytes - init_bytes
          )
          return (message,)
  
      def track_pop(message):
          init_bytes = cs.flatten(message).nbytes
          message, data = codec.pop(message)
          self.pop_bytes[step] += (
              init_bytes - cs.flatten(message).nbytes
          )
          return message, data
  
      return cs.Codec(track_push, track_pop)
  
imports: from abc import ABC, abstractmethod
  from typing import List, NamedTuple, Optional, Tuple, Union, overload
  import craystack
  import numpy
  import torch
  from craystack import codecs
  from craystack.bb_ans import BBANS
  from improved_diffusion.gaussian_diffusion import GaussianDiffusion
label_desc: Use image node_id 14 for calling the BitsBackCodec obj's underlying member method code with example usage: obj._track_codec(codec, step) and returns: cs, codec, message, message,  data