code: def encode(
      messages: ndarray,
      codec: CrayCodec,
      tail_capacity: Optional[int] = None,
      start_buffers: Optional[Sequence[ndarray]] = None,
      tail_capacity_mult: float = 1.5,
  ) -> Tuple[Sequence[ndarray], Sequence[ndarray]]:
      """
      rANS encoder using craystack API.
  
      This function encodes a batched array of messages into bytes via batched
      and interleaved entropy coding with the rANS algorithm. This function can
      also be used to add information to an existing compression buffer.
  
      * ``messages`` should be a JAX-Numpy array of shape
        ``(batch_size, message_length, interleave_level)``. The computation will
        be parallelized over both batches and interleave levels. Each message
        index of the ``interleave_level`` dimension will be coded into the same
        batch output dimension - i.e., the result is a ``batch_size``-length list
        of arrays.
  
      This rANS variant is based on Craystack, which is described in `A tutorial
      on the range variant of asymmetric numeral systems (J. Townsend)
      <https://arxiv.org/abs/2001.09186>`_.
  
      The interleaved functionaliy is described in `Interleaved entropy coders
      (F. Giesen) <https://arxiv.org/abs/1402.3392>`_.
  
      Args:
          messages: A JAX array of messages to be encoded of shape
              ``(batch_size, message_length, interleave_level)``. Note that
              messages in the ``interleave_level`` dimension will be encoded into
              the same batch index.
          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.
          tail_capacity: An optional manual specification for size of the stack
              tail.
          start_buffers: If provided, the function will append to these existing
              compressed buffers rather than creating a new buffer.
          tail_capacity_mult: A multiplication factor used for calculating the
              size of the tail based on the input message size. Compression
              will silently fail if this value is too small. However, larger
              values imply in longer compression/decompression times.
  
      Returns:
          A 2-tuple containing:
              A list of arrays of compressed data, one for each batch element.
              The final CDF state.
      """
      if tail_capacity is None:
          tail_capacity = int(tail_capacity_mult * messages[0].size)
  
      if start_buffers is None:
          message_buffers = jnp.stack(
              [
                  craymessage_to_array(
                      empty_message(messages.shape[2:], tail_capacity)
                  )[0]
                  for _ in range(messages.shape[0])
              ]
          )
          tail_limits = jnp.stack(
              [tail_capacity for _ in range(messages.shape[0])]
          )
      else:
          message_buffers, tail_limits = convert_to_embedded(
              start_buffers, messages.shape[2:], tail_capacity
          )
          tail_limits -= 1  # make sure we don't overwrite info
  
      def vmap_fun(message, message_buffer, tail_limit, init_cdf_state):
          return _jax_encode(
              message,
              array_to_craymessage(
                  message_buffer, messages.shape[2:], tail_limit
              ),
              codec,
              init_cdf_state,
          )
  
      map_fun = jax.vmap(vmap_fun)
  
      (compressed_messages, lims), cdf_state = map_fun(
          messages, message_buffers, tail_limits, codec.cdf_state
      )
  
      head_size = int(jnp.prod(jnp.array(messages.shape[2:])) * 4)
      array_list = []
      for i, (cm, lim) in enumerate(zip(compressed_messages, lims)):
          if lim - head_size < 0:
              raise RuntimeError(
                  f"tail_capacity at batch idx {i} not sufficient for coding."
              )
  
          splits = jnp.split(cm, [head_size, lim - 1])
          array_list.append(jnp.concatenate((splits[0], splits[2])))
  
      return array_list, cdf_state
  
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 3 for calling a global function with example usage: encode(messages, codec, tail_capacity, start_buffers, tail_capacity_mult) and returns: array_list,  cdf_state, _jax_encode