code: def _jax_decode(
      message_len: ndarray,
      int_array: ndarray,
      inverse_cdf_fun: Callable,
      init_inverse_cdf_state: Any,
  ) -> ndarray:
      """JAX decoding backend function."""
      idx = jnp.array(0, jnp.int32)
      bit_idx = jnp.array(0, jnp.int32)
      message = jnp.zeros(message_len, dtype=jnp.uint8)
      value = jnp.array(0, jnp.uint32)
      high = jnp.array(MAX_CODE, dtype=jnp.uint32)
      low = jnp.array(0, dtype=jnp.uint32)
      precision = jnp.array(16, dtype=jnp.uint32)
  
      def dequeue_one_symbol(msg_idx, dequeue_state):
          message = dequeue_state[0]
          (
              value,
              symbol,
              idx,
              bit_idx,
              high,
              low,
              inverse_cdf_state,
          ) = _dequeue(
              *dequeue_state[1:7], inverse_cdf_fun, *dequeue_state[7:]
          )
  
          message = jax.ops.index_update(message, msg_idx, symbol)
  
          return (
              message,
              value,
              int_array,
              idx,
              bit_idx,
              high,
              low,
              inverse_cdf_state,
              precision,
          )
  
      # read out one bit and append it to our byte
      def append_bit_to_value(loop_idx, vals):
          value, idx, bit_idx, int_array = vals
  
          bit, idx, bit_idx = _retrieve_bit(idx, bit_idx, int_array)
  
          value <<= 1
          value += bit
  
          return value, idx, bit_idx, int_array
  
      value, idx, bit_idx, _ = lax.fori_loop(
          0,
          jnp.int64(precision * 2),
          append_bit_to_value,
          (value, idx, bit_idx, int_array),
      )
  
      # extract all symbols from the compressed array
      message = lax.fori_loop(
          0,
          message_len,
          dequeue_one_symbol,
          (
              message,
              value,
              int_array,
              idx,
              bit_idx,
              high,
              low,
              init_inverse_cdf_state,
              precision,
          ),
      )[0]
  
      return message
  
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 7 for calling a global function with example usage: _jax_decode(message_len, int_array, inverse_cdf_fun, init_inverse_cdf_state) and returns: message, message,  value,  int_array,  idx,  bit_idx,  high,  low,  inverse_cdf_state,  precision, value,  idx,  bit_idx,  int_array