code: def bitsback_ans_codec(
      latent_prior_codec: CrayCodec,
      latent_posterior_codec_maker: Callable[[ndarray], CrayCodec],
      obs_codec_maker: Callable[[ndarray], CrayCodec],
      latent_shape: Sequence[int],
      message_dtype: str,
  ) -> CrayCodec:
      r"""
      Bitsback with Asymmetric Numeral Systems (BB-ANS) craystack codec.
  
      This function implements BB-ANS as described in 'Practical Lossless
      Compression with Latent Variables using Bits Back Coding', Townsend
      et. al. 2019. Please cite the original paper where appropriate:
      https://arxiv.org/abs/1901.04866
  
      When compressing a symbol :math:`x` with a latent variable model
      (e.g. VAE), we might not have access to the evidence :math:`p(x)`, only
      latent prior :math:`p(z)`, latent approximate posterior :math:`q(z|x)`,
      and observational contitional likelihood :math:`p(x|z)`. BB-ANS allows
      us to compress with these auxiliary distributions, while achieving a rate
      close to that of compressing with :math:`p(x)`, depending on the quality
      of the approximate posterior :math:`q(z|x)`.
  
      The encoding scheme for a symbol :math:`x` is
          1. decode :math:`z` with :math:`q(z|x)`
          2. encode :math:`x` with :math:`p(x|z)`
          3. encode :math:`z` with :math:`p(z)`
  
      Decoding proceeds in reverse order
          3. decode :math:`z` with :math:`p(z)`
          2. decode :math:`x` with :math:`p(x|z)`
          1. encode :math:`z` with :math:`q(z|x)` <-- bits-back step
  
      Step 1. decreases the ANS state by approximately :math:`\log_2 q(z|x)`,
      while the other 2 steps together increase it by :math:`-\log_2 p(z)p(x|z)`,
      implying the rate is :math:`ELBO(x) = -\log_2 \frac{p(z)p(x|z)}{q(z|x)}`.
      When :math:`q(z|x) = p(z|x)`, i.e. the approximate posterior perfectly
      matches the true posterior, then :math:`ELBO(x) = -\log_2 p(x)`.
  
      Popping the initial :math`z`'s requires that there be information on the
      ANS stack. This is known as the 'initial bits problem'. The CrayTail
      component of CrayCompressedMessage can automatically return random bits if
      the stack is depleted. For more details, see the ``pop_symbols`` method in
      file ``neuralcompression.entropy_coders.craystack._backend.py``
  
      Args:
          latent_posterior_codec_maker: A ``Callable`` that receives the symbols
              to be encoded and returns a ``CrayCodec`` representing ``q(z|x)``
              used in step 1.
          obs_codec_maker: A ``Callable``that receives the latents and returns
              a ``CrayCodec` representing ``p(x|z)`` used in step 2.
          latent_prior_codec: A ``CrayCodec`` representing ``p(z)`` in step 3.
          latent_shape: A ``Sequence[int]``, representing the shape of the latents
              ``z``, used to construct a temporary message that holds ``z``
              during encoding and decoding.
          message_dtype: Data type for input and output messages. Muster be the
              same message_dtype as the codecs returned by ``obs_codec_maker``.
  
      Returns:
          A Craystack codec that can be input to the coder.
      """
  
      def push(symbols, compressed_message, cdf_state):
          # pop latents with approximate posterior q(latents|symbols)
          latent_posterior_codec = latent_posterior_codec_maker(symbols)
          compressed_message, latents, _ = latent_posterior_codec.pop(
              0,
              compressed_message,
              jnp.zeros(
                  latent_shape,
                  dtype=latent_posterior_codec.message_dtype,
              ),
              latent_posterior_codec.cdf_state,
              latent_posterior_codec.allow_empty_pops,
          )
  
          # push symbols with conditional likelihood p(symbols|latents)
          obs_codec = obs_codec_maker(latents)
          compressed_message, _ = obs_codec.push(
              symbols, compressed_message, obs_codec.cdf_state
          )
  
          # push latents with prior p(latents)
          compressed_message, _ = latent_prior_codec.push(
              latents,
              compressed_message,
              latent_prior_codec.cdf_state,
          )
  
          return compressed_message, None
  
      def pop(msg_index, compressed_message, message, cdf_state):
          # pop latents with prior p(latents)
          compressed_message, latents, _ = latent_prior_codec.pop(
              0,
              compressed_message,
              jnp.zeros(
                  latent_shape, dtype=latent_prior_codec.message_dtype
              ),
              latent_prior_codec.cdf_state,
              latent_prior_codec.allow_empty_pops,
          )
  
          # pop symbols with conditional likelihood p(symbols|latents)
          obs_codec = obs_codec_maker(latents)
          compressed_message, message, _ = obs_codec.pop(
              msg_index,
              compressed_message,
              message,
              obs_codec.cdf_state,
              obs_codec.allow_empty_pops,
          )
          symbols = message[msg_index]
  
          # push latents with approximate posterior q(latents|symbols)
          # this is the bits-back step!
          latent_posterior_codec = latent_posterior_codec_maker(symbols)
          compressed_message, _ = latent_posterior_codec.push(
              latents,
              compressed_message,
              latent_posterior_codec.cdf_state,
          )
  
          return compressed_message, message, None
  
      return CrayCodec(push, pop, None, message_dtype)
  
file_path: NeuralCompression/projects/jax_entropy_coders/craystack/codecs.py
imports: from typing import Callable, NamedTuple, Optional, Sequence, Tuple
  import jax.lax
  import jax.numpy
  from jax.numpy import ndarray
  from ._backend import CrayCompressedMessage, peek, pop_symbols, push_symbols
label_desc: Use image node_id 3 for calling a global function with example usage: bitsback_ans_codec(latent_prior_codec, latent_posterior_codec_maker, obs_codec_maker, latent_shape, message_dtype) and returns: CrayCodec, compressed_message,  None, compressed_message,  message,  None