code: def _append_bit(
      bit: bool, int_array: ndarray, byte: int, idx: int, bit_idx: int
  ) -> Tuple[ndarray, int, int, int]:
      """Append a single bit to the integer bit array."""
  
      def true_branch(vals):
          return jax.ops.index_update(
              vals[0], vals[1], vals[2]
          ), jnp.uint8(0)
  
      def false_branch(vals):
          return vals[0], vals[2]
  
      byte <<= 1
      byte |= bit
  
      bit_idx += 1
      int_array, byte = lax.cond(
          bit_idx == 8,
          true_branch,
          false_branch,
          (int_array, idx, byte),
      )
      idx += bit_idx // 8
      bit_idx = jnp.mod(bit_idx, 8)
  
      return int_array, byte, idx, bit_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 1 for calling a global function with example usage: _append_bit(bit, int_array, byte, idx, bit_idx) and returns: int_array,  byte,  idx,  bit_idx, jax,  jnp, vals,  vals