code: def encode(
      messages: ndarray,
      cdf_fun: Callable,
      init_cdf_states: Any,
      buffer_oversize: float = 1.5,
  ) -> List[bytes]:
      """
      Encode a batch of messages via sequential arithmetic coding.
  
      This function encodes a batch of messages into a list of byte arrays using
      cumulative distribution Python functionals. Its unique advantage is that
      the CDF can be specified functionally rather than via look-up table, which
      allows use of complicated functionals such as neural networks.
  
      * ``messages`` should be a JAX-Numpy array of shape
        ``(batch_size, message_length)``. The computation will be parallelized
        over batches.
      * ``cdf_fun`` should be a callable of the form
        ``f(symbol, cdf_state) -> (cdf_low, cdf_high, cdf_state)``, where
        ``cdf_low`` is :math:`P(x <= symbol)` and `cdf_high` is
        :math:`P(x <= symbol+1)`. ``cdf_state`` can be any collection of JAX
        arrays that is used internally by ``cdf_fun`` during coding. The CDF
        probabilities should be on an integer scale between ``0`` and
        ``2 ** precision``, where precision is 16 (for a 16-bit working integer).
  
      Args:
          messages: An input JAX ndarray where the first dimension is the batch
              dimension.
          cdf_fun: A Python callable for calculating symbol probabilities.
          init_cdf_state: The initialization state of the CDF function (contains
              CDF array or can be used for conditional probabilites).
          buffer_oversize: The compression buffer size is
              `buffer_oversize*len(messages[0])`.
  
      Returns:
          A list of length ``batch_size`` where each element in the list is a
              compressed byte array of the content in that batch index.
      """
  
      def run_encode(message, init_cdf_state):
          return _single_batch_encode(
              message, cdf_fun, init_cdf_state, buffer_oversize
          )
  
      map_fun = jax.vmap(run_encode)
  
      bits, num_bytes = map_fun(messages, init_cdf_states)
  
      if any([i > len(b) for b, i in zip(bits, num_bytes)]):
          raise RuntimeError(
              "Message buffer overrun - try increasing buffer_oversize."
          )
  
      bits = [bytes(b[: i + 1]) for b, i in zip(bits, num_bytes)]
  
      return bits
  
file_path: NeuralCompression/projects/jax_entropy_coders/jax_arithmetic_coder.py
imports: from typing import Any, Callable, List, Tuple
  import jax
  import jax.lax
  import jax.numpy
  from jax.numpy import ndarray
label_desc: Use image node_id 10 for calling a global function with example usage: encode(messages, cdf_fun, init_cdf_states, buffer_oversize) and returns: bits, _single_batch_encode