code: def _dequeue(
      value: ndarray,
      int_array: ndarray,
      idx: ndarray,
      bit_idx: ndarray,
      high: ndarray,
      low: ndarray,
      inverse_cdf_fun: Callable,
      inverse_cdf_state: Any,
      precision: ndarray,
  ) -> Tuple[ndarray, Any, ndarray, ndarray, ndarray, ndarray, Any]:
      """
      Dequeue a single symbol from the compressed ``int_array``.
  
      Args:
          value: The current location within the CDF range used to decode symbol.
          int_array: A 8-bit integer array holding the so-far compressed message.
          idx: The index of the operating location within `int_array`.
          bit_idx: The index of the current bit within the operating byte.
          high: The current high of the range.
          low: The current low of the range.
          inverse_cdf_fun: A callable that can return a symbol, its low and high
              probabilities, and an updated `inverse_cdf_state` given a value and
              the current ``inverse_cdf_state``.
          inverse_cdf_state: The current state of the CDF function (contains CDF
              array or can be used for conditional probabilites).
          precision: An integer specifying the operating precision of the
              algorithm and max CDF value.
  
      Returns:
          A tuple containing updated versions of ``value``, ``high``, ``low``,
              ``idx``, ``bit_idx``, ``int_array``.
      """
  
      # jax functions - these will all be jit-compiled for our while loop
      def high_cross_half(vals):
          """No need to do anything, bit is 0."""
          return vals
  
      def low_cross_half(vals):
          """Based on low increasing, adjust limits."""
          value, high, low = vals
  
          value -= ONE_HALF
          low -= ONE_HALF
          high -= ONE_HALF
  
          return value, high, low
  
      def cross_fourth(vals):
          """Range got too narrow, open up."""
          value, high, low = vals
  
          value -= ONE_FOURTH
          low -= ONE_FOURTH
          high -= ONE_FOURTH
  
          return value, high, low
  
      branches = [high_cross_half, low_cross_half, cross_fourth]
  
      def choose_branch_and_execute(vals):
          """if/else bad for jax, so we use a switch with argmin instead."""
          value, high, low, idx, bit_idx, int_array = vals
  
          # this will return true for the first one of these that is true
          branch_idx = lax.argmax(
              jnp.array(
                  [
                      (high < ONE_HALF),
                      (low >= ONE_HALF),
                      (low >= ONE_FOURTH) * (high < THREE_FOURTHS),
                  ]
              ),
              0,
              jnp.int32,
          )
  
          # execute chosen branch
          value, high, low = lax.switch(
              branch_idx, branches, (value, high, low)
          )
  
          low <<= 1
          high <<= 1
          high += 1
          value <<= 1
          bit, idx, bit_idx = _retrieve_bit(idx, bit_idx, int_array)
          value += bit * (idx < len(int_array))
  
          return value, high, low, idx, bit_idx, int_array
  
      def condition(vals):
          """Once our range has converged, we break the while loop."""
          return (
              (vals[1] < ONE_HALF)
              + (vals[2] >= ONE_HALF)
              + (vals[2] >= ONE_FOURTH) * (vals[1] < THREE_FOURTHS)
          ) > 0
  
      # calculate our current range window
      span = jnp.uint64(high) - jnp.uint64(low) + 1
  
      scaled_value = jnp.uint16(
          (
              (jnp.uint64(value) - jnp.uint64(low) + 1)
              * jnp.uint64(PREC_COUNT)
              - 1
          )
          / span
      )
  
      # get the low and high points of the CDF for current symbol
      (
          range_low,
          range_high,
          symbol,
          inverse_cdf_state,
      ) = inverse_cdf_fun(scaled_value, inverse_cdf_state)
  
      # narrow the range using cdf values
      high = (
          low
          + jnp.uint32(
              (span * jnp.uint64(range_high)) >> jnp.uint64(precision)
          )
          - 1
      )
      low = low + jnp.uint32(
          (span * jnp.uint64(range_low)) >> jnp.uint64(precision)
      )
  
      # build and execute while loop
      value, high, low, idx, bit_idx, _ = lax.while_loop(
          condition,
          choose_branch_and_execute,
          (value, high, low, idx, bit_idx, int_array),
      )
  
      return value, symbol, idx, bit_idx, high, low, inverse_cdf_state
  
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 5 for calling a global function with example usage: _dequeue(value, int_array, idx, bit_idx, high, low, inverse_cdf_fun, inverse_cdf_state, precision) and returns: value,  symbol,  idx,  bit_idx,  high,  low,  inverse_cdf_state, vals, value,  high,  low, value,  high,  low, value,  high,  low,  idx,  bit_idx,  int_array, unknown