code: def _single_batch_decode(
      data: ndarray,
      message_len: ndarray,
      inverse_cdf_fun: Callable,
      init_inverse_cdf_state: Any,
  ) -> jnp.ndarray:
      """Decodes one index of a batch."""
      message = _jax_decode(
          message_len=message_len,
          int_array=data,
          inverse_cdf_fun=inverse_cdf_fun,
          init_inverse_cdf_state=init_inverse_cdf_state,
      )
  
      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 9 for calling a global function with example usage: _single_batch_decode(data, message_len, inverse_cdf_fun, init_inverse_cdf_state) and returns: message