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: Posterior codec for BB-ANS.
  
  Args:
      discr_data: The discrete data in [0, 2 ** self.data_prec - 1],
          on which the posterior is conditioned.
  Returns:
      The craystack codec corresponding to the posterior
      of the Bits-Back codec.
docstr_tok: ["Posterior","codec","for","BB-ANS",".","Args",":","discr_data",":","The","discrete","data","in","[","0",",","2","*","*","self.data_prec","-","1","]",",","on","which","the","posterior","is","conditioned",".","Returns",":","The","craystack","codec","corresponding","to","the","posterior","of","the","Bits-Back","codec","."]
code: def posterior(self, discr_data: np.ndarray) -> cs.Codec:
      """
      Posterior codec for BB-ANS.
  
      Args:
          discr_data: The discrete data in [0, 2 ** self.data_prec - 1],
              on which the posterior is conditioned.
      Returns:
          The craystack codec corresponding to the posterior
          of the Bits-Back codec.
      """
      data = self.data_min + discr_data.astype(float) * (
          self.data_max - self.data_min
      ) / ((1 << self.data_prec) - 1)
      top_post_params = self.model.infer_fn(data)
  
      def posterior_push(
          message: tuple, discr_latents: List[DiscrLatent]
      ) -> tuple:
          """
          Push indices idx_1, ..., idx_{steps-1} corresponding to the
          latent variables z_1, ..., z_{steps-1} using the Gaussian posteriors
          q(z_{step}|z_{step+1}, x) or q(z_{steps-1}|discr_data) and
          discretization given by bins that have equal mass under Gaussian
          priors with mean/stdd given by the parameters of the latents.
  
          Args:
              message: The craystack message to push to.
              discr_latents: The list of (discretized) latents
                  to be pushed to the message.
          Returns:
              The new craystack message with the indices appended.
          """
          assert len(discr_latents) == self.model.steps
          for step in range(1, self.model.steps + 1):
              # first get the posterior params
              if step < self.model.steps:
                  next_latent = self._to_latent(discr_latents[step])
                  post_params = self.model.generate_fn(
                      next_latent, step=step, data=data
                  )
              else:
                  post_params = top_post_params
  
              discr_latent = discr_latents[step - 1]
              codec = self._track_codec(
                  self._post_codec(post_params, discr_latent.params),
                  step=step,
              )
              (message,) = codec.push(message, discr_latent.inds)
          return (message,)
  
      def posterior_pop(
          message: tuple,
      ) -> Tuple[tuple, List[DiscrLatent]]:
          """
          Pop indices idx_1, ..., idx_{steps-1} corresponding to the
          latent variables z_1, ..., z_{steps-1} using the Gaussian posteriors
          q(z_{step}|z_{step+1}, discr_data) or q(z_{steps-1}|discr_data)
          and discretization given by bins that have equal mass
          under Gaussian priors p(z_{step}|z_{step+1}) or p(z_{steps-1}).
  
          Args:
              message: The craystack message to pop from.
          Returns:
              A tuple containing the new craystack message
              and the popped indices.
          """
          discr_latents: List[DiscrLatent] = []
          for step in range(self.model.steps, 0, -1):
              if step == self.model.steps:
                  prior_params = GaussianParams(
                      mean=np.zeros(self.data_shape),
                      stdd=np.ones(self.data_shape),
                  )
                  post_params = top_post_params
              else:
                  latent = self._to_latent(discr_latents[-1])
                  prior_params = self.model.generate_fn(
                      latent, step=step
                  )
                  post_params = self.model.generate_fn(
                      latent, step=step, data=data
                  )
  
              codec = self._post_codec(post_params, prior_params)
              message, latent_inds = self._track_codec(
                  codec, step=step
              ).pop(message)
              discr_latents.append(
                  DiscrLatent(inds=latent_inds, params=prior_params)
              )
  
          if self.data_count == 0:
              # compute extra bytes needed for encoding
              self.extra_bytes = (
                  self.initial_bytes - cs.flatten(message).nbytes
              )
  
          return message, discr_latents[::-1]
  
      return cs.Codec(posterior_push, posterior_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 12 for calling the BitsBackCodec obj's underlying member method code with example usage: obj.posterior(discr_data) and returns: cs, message, message,  discr_latents