code: def default_rans_codec(
      cdf_fun: Callable[
          [ndarray, Sequence[ndarray]],
          Tuple[ndarray, ndarray, Sequence[ndarray]],
      ],
      inverse_cdf_fun: Callable[
          [ndarray, Sequence[ndarray]],
          Tuple[ndarray, ndarray, ndarray, Sequence[ndarray]],
      ],
      init_cdf_state: Sequence[ndarray],
      precision: int = 16,
      message_dtype: str = "int64",
      allow_empty_pops: bool = False,
  ) -> CrayCodec:
      """
      Default rANS codec for Craystack.
  
      This function takes as input a functional CDF, an inverse CDF and an
      initial CDF state and converts to a Craystack codec that can be input by
      defining push and pop operations that can be used by the internals of the
      rANS algorithm.
  
      Args:
          cdf_fun: A ``Callable`` that maps a symbol and a ``cdf_state`` to a
              3-length tuple containing 1) the low CDF value for the symbol, 2)
              the high CDF value for the symbol and 3) the updated CDF state.
          inverse_cdf_fun: A ``Callable`` that maps a value and a ``cdf_state``
              to a 4-length tuple containing 1) the low CDF value for the
              symbol, 2) the high value for the symbol, 3) the symbol located at
              the given value, and 4) an updated CDF state.
          precision: An integer specifying codec operating precision.
          message_dtype: Data type for input and output messages.
  
      Returns:
          A Craystack codec that can be input to the coder.
      """
  
      def push(symbols, compressed_message, cdf_state):
          cdf_low, cdf_high, cdf_state = cdf_fun(symbols, cdf_state)
          return (
              push_symbols(
                  compressed_message, cdf_low, cdf_high, precision
              ),
              cdf_state,
          )
  
      def pop(
          msg_index,
          compressed_message,
          message,
          cdf_state,
          allow_empty_pops=False,
      ):
          cdf_value = peek(compressed_message, precision)
          cdf_low, cdf_high, symbols, cdf_state = inverse_cdf_fun(
              cdf_value, cdf_state
          )
          return (
              pop_symbols(
                  compressed_message,
                  cdf_value,
                  cdf_low,
                  cdf_high,
                  precision,
                  allow_empty_pops,
              ),
              lax.dynamic_update_index_in_dim(
                  message, symbols.astype(message_dtype), msg_index, 0
              ),
              cdf_state,
          )
  
      return CrayCodec(
          push, pop, init_cdf_state, message_dtype, allow_empty_pops
      )
  
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 2 for calling a global function with example usage: default_rans_codec(cdf_fun, inverse_cdf_fun, init_cdf_state, precision, message_dtype, allow_empty_pops) and returns: CrayCodec, push_symbols,  cdf_state, pop_symbols,  lax,  cdf_state