code: def _single_batch_encode(
      message: jnp.ndarray,
      cdf_fun: Callable,
      init_cdf_state: Any,
      buffer_oversize: float,
  ) -> Tuple[ndarray, int]:
      """Encodes one index of a batch."""
      bits, idx = _jax_encode(
          message, cdf_fun, init_cdf_state, buffer_oversize
      )
  
      return bits, 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 8 for calling a global function with example usage: _single_batch_encode(message, cdf_fun, init_cdf_state, buffer_oversize) and returns: bits,  idx