code: def fixed_array_cdf_codec(
      cdf_array: ndarray,
      precision: int = 16,
      message_dtype: str = "int64",
      allow_empty_pops: bool = False,
  ) -> CrayCodec:
      """
      Simple Craystack codec using CDF array.
  
      Given an ``N+1``-length array that defines a cumulative distribution
      function for ``N`` symbols, this function builds a Craystack codec by
      defining push and pop operations that can be used by the internals of the
      rANS algorithm.
  
      Args:
          cdf_array: An ``N+1``-length array defining the CDF for ``N`` symbols
              where the frequency count of symbol ``s`` is
              ``cdf_array[s+1] - cdf_array[s]``.
          precision: An integer specifying codec operating precision.
          message_dtype: Data type for input and output messages.
          allows_empty_pops: A boolean that specifys if CrayTail should be
              allowed to create random bits on demand. See the ``pop_symbols``
              method in ``entropy_coders.craystack._backend.py``.
  
      Returns:
          A Craystack codec that can be input to the coder.
      """
      init_cdf_state = (cdf_array,)
  
      def _lookup(cdf, symbols):
          return jnp.take_along_axis(cdf, symbols[..., None], -1)[
              ..., 0
          ]
  
      def cdf_fun(symbols, cdf_state):
          cdf = cdf_state[0]
          return (
              _lookup(cdf, symbols),
              _lookup(cdf, symbols + 1),
              cdf_state,
          )
  
      def inverse_cdf_fun(cdf_value, cdf_state):
          cdf = cdf_state[0]
          symbols = (
              jnp.argmin(jnp.expand_dims(cdf_value, -1) >= cdf, axis=-1)
              - 1
          )
          return (
              _lookup(cdf, symbols),
              _lookup(cdf, symbols + 1),
              symbols,
              cdf_state,
          )
  
      return default_rans_codec(
          cdf_fun,
          inverse_cdf_fun,
          init_cdf_state,
          precision,
          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 1 for calling a global function with example usage: fixed_array_cdf_codec(cdf_array, precision, message_dtype, allow_empty_pops) and returns: default_rans_codec, jnp, _lookup,  _lookup,  cdf_state, _lookup,  _lookup,  symbols,  cdf_state