code: def _jax_decode(
      compressed_message: ndarray,
      tail_limit: int,
      message_len: int,
      message_shape: Sequence[int],
      codec: CrayCodec,
      cdf_state: Sequence[ndarray],
  ) -> Tuple[Tuple[ndarray, int], ndarray, Sequence[ndarray]]:
      """
      JAX rANS decoding function.
  
      At a high level, this function takes a stack of information
      (``compressed_message``) and peeks at the top of the stack to see what the
      current symbol is. After identifying the symbol, this function pops a
      number of bits from the top of the stack approximately equal to the
      information content of the symbol (i.e. ``-log(symbol probability)``). This
      is done ``message_len`` times until the full message is retrieved.
  
      Args:
          compressed_message: The input stack containing the compressed meessage.
          tail_limit: A pointer to the current end of the tail.
          message_len: The size of the message to be decoded.
          message_shape: The message shape containing the interleaved dimension
              size.
          codec: A named tuple object containing functions for push and pop
              operations, as well as an initial state fo the CDF functions (for
              context-adaptive coding) and a data type specification for the
              message.
          cdf_state: The initialization state of the inverse CDF function
              (contains CDF array or can be used for conditional probabilites).
  
      Returns:
          A 3-tuple containing:
              The decoded messages of size
                      ``(message_len, *message_shape)``.
              A byte array of compressed data after removing the target
                  message.
              The final CDF state.
      """
      message = jnp.zeros(
          (message_len, *message_shape), dtype=codec.message_dtype
      )
  
      def pop_one_symbol(msg_index, vals):
          return codec.pop(msg_index, *vals)
  
      result = lax.fori_loop(
          0,
          message_len,
          pop_one_symbol,
          (
              array_to_craymessage(
                  compressed_message, message_shape, tail_limit
              ),
              message,
              cdf_state,
          ),
      )
  
      return craymessage_to_array(result[0]), result[1], result[2]
  
file_path: NeuralCompression/projects/jax_entropy_coders/craystack/coder.py
imports: from typing import Optional, Sequence, Tuple
  import jax
  import jax.lax
  import jax.numpy
  from jax.numpy import ndarray
  from ._backend import CrayCompressedMessage, array_to_craymessage, convert_to_embedded, craymessage_to_array, empty_message
  from .codecs import CrayCodec
label_desc: Use image node_id 2 for calling a global function with example usage: _jax_decode(compressed_message, tail_limit, message_len, message_shape, codec, cdf_state) and returns: craymessage_to_array,  result,  result, codec