code: def _jax_encode(
      message: ndarray,
      compressed_message: CrayCompressedMessage,
      codec: CrayCodec,
      cdf_state: Sequence[ndarray],
  ) -> Tuple[Tuple[ndarray, int], Sequence[ndarray]]:
      """
      JAX rANS encoding function.
  
      At a high level, this function takes a message and appends the information
      content of that message to a byte stack array. The size and content of the
      append operation depends on the probability of each symbol in the message
      occurring.
  
      Args:
          message: The array to be encoded.
          compressed_message: A stack byte array to append the compressed symbols
              to.
          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 CDF function
              (contains CDF array or can be used for conditional probabilites).
  
      Returns:
          A 2-tuple containing:
              A stack with the information content of the compressed symbols.
              The final CDF state.
      """
      message_len = message.shape[0]
  
      def push_one_symbol(msg_index, vals):
          return codec.push(
              message[message_len - msg_index - 1], *vals
          )  # rANS encodes backwards
  
      compressed_message_final, cdf_state_final = lax.fori_loop(
          0,
          message_len,
          push_one_symbol,
          (compressed_message, cdf_state),
      )
  
      return (
          craymessage_to_array(compressed_message_final),
          cdf_state_final,
      )
  
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 1 for calling a global function with example usage: _jax_encode(message, compressed_message, codec, cdf_state) and returns: craymessage_to_array,  cdf_state_final, codec