code: def _jax_encode(
      message: jnp.ndarray,
      cdf_fun: Callable,
      init_cdf_state: Any,
      buffer_oversize: float,
  ) -> Tuple[ndarray, int]:
      """JAX encoding backend function."""
      # array initialization
      idx = jnp.array(0, jnp.int32)
      bit_idx = jnp.array(0, jnp.int32)
      int_array = jnp.zeros(
          int(buffer_oversize * len(message)), dtype=jnp.uint8
      )  # message can grow
      high = jnp.array(MAX_CODE, dtype=jnp.uint32)
      low = jnp.array(0, dtype=jnp.uint32)
      precision = jnp.array(16, dtype=jnp.uint32)
      pending_bits = jnp.array(0, dtype=jnp.int32)
  
      # loop iterate - this enqueues one symbol into the compressed array
      def enqueue_one_symbol(msg_index, enqueue_state):
          symbol = message[msg_index]
          (
              int_array,
              byte,
              idx,
              bit_idx,
              pending_bits,
              high,
              low,
              cdf_state,
          ) = enqueue_state
          return _enqueue(
              symbol,
              int_array,
              byte,
              idx,
              bit_idx,
              pending_bits,
              high,
              low,
              cdf_fun,
              cdf_state,
              precision,
          )
  
      # loop over every symbol in the message and add it to the compressed array
      (
          int_array,
          byte,
          idx,
          bit_idx,
          pending_bits,
          _,
          low,
          _,
      ) = lax.fori_loop(
          0,
          len(message),
          enqueue_one_symbol,
          (
              int_array,
              int_array[0],
              idx,
              bit_idx,
              pending_bits,
              high,
              low,
              init_cdf_state,
          ),
      )
  
      # append the final trailing bits
      pending_bits += 1
  
      def append_final_0(final_append):
          pending_bits, int_array, byte, idx, bit_idx = final_append
          int_array, byte, idx, bit_idx = _append_bit_and_pending(
              False, pending_bits, int_array, byte, idx, bit_idx
          )
          return int_array, byte, idx, bit_idx
  
      def append_final_1(final_append):
          pending_bits, int_array, byte, idx, bit_idx = final_append
          int_array, byte, idx, bit_idx = _append_bit_and_pending(
              True, pending_bits, int_array, byte, idx, bit_idx
          )
          return int_array, byte, idx, bit_idx
  
      int_array, byte, idx, bit_idx = lax.cond(
          low < ONE_FOURTH,
          append_final_0,
          append_final_1,
          (pending_bits, int_array, byte, idx, bit_idx),
      )
  
      int_array = jax.ops.index_update(
          int_array, idx, byte << (8 - bit_idx)
      )
  
      return int_array, idx
  
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 6 for calling a global function with example usage: _jax_encode(message, cdf_fun, init_cdf_state, buffer_oversize) and returns: int_array,  idx, _enqueue, int_array,  byte,  idx,  bit_idx, int_array,  byte,  idx,  bit_idx