code: def _enqueue(
      symbol: Any,
      int_array: ndarray,
      byte: int,
      idx: ndarray,
      bit_idx: ndarray,
      pending_bits: ndarray,
      high: ndarray,
      low: ndarray,
      cdf_fun: Callable,
      cdf_state: Any,
      precision: ndarray,
  ) -> Tuple[
      ndarray, int, ndarray, ndarray, ndarray, ndarray, ndarray, Any
  ]:
      """
      Enqueue a single symbol into the compressed int_array.
  
      Args:
          symbol: A single symbol to code into ``int_array``.
          int_array: A 8-bit integer array holding the so-far compressed message.
          byte: The current operating byte of ``int_array``.
          idx: The index of `byte`'s location within ``int_array``.
          bit_idx: The index of the current bit within ``byte``.
          pending_bits: The number of pending bits to encode (arises from
              convergent CDF).
          high: The current high of the range.
          low: The current low of the range.
          cdf_fun: A callable that can return the probability of the symbol given
              a ``cdf_state``.
          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 ``int_array``, ``byte``,
              ``idx``, ``bit_idx``, ``pending_bits``, ``high``, ``low``, and
              ``cdf_state``.
      """
  
      # jax functions - these will all be jit-compiled for our while loop
      def decrease_high(vals):
          """Based on high decreasing, we append a 0 most significant bit."""
          high, low, pending_bits, int_array, byte, idx, bit_idx = vals
          int_array, byte, idx, bit_idx = _append_bit_and_pending(
              False, pending_bits, int_array, byte, idx, bit_idx
          )
          pending_bits *= 0
  
          return high, low, pending_bits, int_array, byte, idx, bit_idx
  
      def increase_low(vals):
          """Based on low increasing, we append a 1 most significant bit."""
          high, low, pending_bits, int_array, byte, idx, bit_idx = vals
          int_array, byte, idx, bit_idx = _append_bit_and_pending(
              True, pending_bits, int_array, byte, idx, bit_idx
          )
          pending_bits *= 0
  
          return high, low, pending_bits, int_array, byte, idx, bit_idx
  
      def renormalize_range(vals):
          """Range is too low for numerical precision - this fn readjusts."""
          high, low, pending_bits, _, _, _, _ = vals
          pending_bits += 1
          low -= ONE_FOURTH
          high -= ONE_FOURTH
  
          return (high, low, pending_bits) + vals[3:]
  
      branches = [decrease_high, increase_low, renormalize_range]
  
      def choose_branch_and_execute(vals):
          """if/else bad for jax, so we use a switch with argmin instead."""
          high, low, pending_bits, int_array, byte, idx, bit_idx = 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
          (
              high,
              low,
              pending_bits,
              int_array,
              byte,
              idx,
              bit_idx,
          ) = lax.switch(
              branch_idx,
              branches,
              (high, low, pending_bits, int_array, byte, idx, bit_idx),
          )
  
          # perform after-shifts
          high <<= 1
          high += 1
          low <<= 1
          high &= MAX_CODE
          low &= MAX_CODE
  
          return high, low, pending_bits, int_array, byte, idx, bit_idx
  
      def condition(vals):
          """Once our range has converged, we break the while loop."""
          return (
              (vals[0] < ONE_HALF)
              + (vals[1] >= ONE_HALF)
              + (vals[1] >= ONE_FOURTH) * (vals[0] < THREE_FOURTHS)
          ) > 0
  
      # calculate our current range window
      span = jnp.uint64(high) - jnp.uint64(low) + 1
  
      # get the low and high points of the CDF for current symbol
      range_low, range_high, cdf_state = cdf_fun(symbol, 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
      (
          high,
          low,
          pending_bits,
          int_array,
          byte,
          idx,
          bit_idx,
      ) = lax.while_loop(
          condition,
          choose_branch_and_execute,
          (high, low, pending_bits, int_array, byte, idx, bit_idx),
      )
  
      return (
          int_array,
          byte,
          idx,
          bit_idx,
          pending_bits,
          high,
          low,
          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 4 for calling a global function with example usage: _enqueue(symbol, int_array, byte, idx, bit_idx, pending_bits, high, low, cdf_fun, cdf_state, precision) and returns: int_array,  byte,  idx,  bit_idx,  pending_bits,  high,  low,  cdf_state, high,  low,  pending_bits,  int_array,  byte,  idx,  bit_idx, high,  low,  pending_bits,  int_array,  byte,  idx,  bit_idx, unknown, high,  low,  pending_bits,  int_array,  byte,  idx,  bit_idx, unknown