code: def decode(
      compressed_bytes: Sequence[ndarray],
      message_len: int,
      message_shape: Sequence[int],
      codec: CrayCodec,
      tail_capacity: Optional[int] = None,
      tail_capacity_mult: float = 1.5,
  ) -> Tuple[ndarray, Sequence[ndarray], Sequence[ndarray]]:
      """
      rANS decoder using craystack API.
  
      This function decodes a batched array of messages into bytes via batched
      and interleaved entropy coding with the rANS algorithm. The function
      returns the compressed buffer at the end to allow the removal of subsequent
      messages.
  
      * ``compressed_bytes`` should be a list of ndarrays representing compressed
        data.
      * ``message_shape`` is a tuple specifying the shape of the decoded output.
        For example, if you compressed a ``(batch_size, message_len,
        interleave_level)`` input, then ``message_shape`` should be a 1-tuple
        containing ``interleave_level``.
  
      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:
          compressed_bytes: A list of arrays with compressed messages of length
              ``batch_size``.
          message_len: Length of message to be decoded.
          message_shape: Size of message - must include interleaved level. E.g.,
              if you encoded a message of size ``(5, 100, 2)``, this should be
              ``(2,)``.
          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 calcualting the
              size of the tail based on the input message size.
  
      Returns:
          A 3-tuple containing:
              The decoded messages of size
                  ``(batch_size, message_len, *message_shape)``.
              A list of arrays of compressed data, one for each batch element
                  (after removing the target message).
              The final CDF state.
      """
      if tail_capacity is None:
          tail_capacity = int(
              tail_capacity_mult
              * message_len
              * jnp.prod(jnp.array(message_shape))
          )
  
      compressed_messages, tail_limits = convert_to_embedded(
          compressed_bytes, message_shape, tail_capacity
      )
  
      def vmap_fun(compressed_message, tail_limit, init_cdf_state):
          return _jax_decode(
              compressed_message,
              tail_limit,
              message_len,
              message_shape,
              codec,
              init_cdf_state,
          )
  
      map_fun = jax.vmap(vmap_fun)
  
      (compressed_messages, lims), messages, cdf_state = map_fun(
          compressed_messages, tail_limits, codec.cdf_state
      )
      lims += 1  # correction factor for avoiding encode overwrites
  
      head_size = int(jnp.prod(jnp.array(message_shape)) * 4)
      array_list = []
      for cm, lim in zip(compressed_messages, lims):
          splits = jnp.split(cm, [head_size, lim - 1])
          array_list.append(jnp.concatenate((splits[0], splits[2])))
  
      return messages, 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 4 for calling a global function with example usage: decode(compressed_bytes, message_len, message_shape, codec, tail_capacity, tail_capacity_mult) and returns: messages,  array_list,  cdf_state, _jax_decode