code: def test_arithmetic_coder_adaptive(shape, alphabet_size):
      # Assumes no knowledge of the source distribution. Instead,
      # adaptively estimates the PDF based on the frequency count of
      # previously encoded symbols.
      batch_size = shape[0]
      message_len = shape[1]
      seed = 7 * batch_size
      rng = np.random.default_rng(seed)
  
      messages = jnp.array(
          rng.integers(
              low=0, high=alphabet_size, size=shape, dtype=np.uint8
          )
      )
  
      cdf_state = (
          jnp.array(
              rng.integers(
                  low=1, high=42, size=(batch_size, alphabet_size)
              )
          ),
      )
  
      def cdf_fun(symbol, cdf_state):
          freqs = cdf_state[0]
          cdf = freqs_to_cdf(freqs)
          return (
              cdf[symbol],
              cdf[symbol + 1],
              (freqs.at[symbol].add(1),),
          )
  
      def inverse_cdf_fun(value, cdf_state):
          freqs = cdf_state[0]
          cdf = freqs_to_cdf(freqs)
          symbol = jnp.argmin(value >= cdf) - 1
          return (
              cdf[symbol],
              cdf[symbol + 1],
              symbol,
              (freqs.at[symbol].add(1),),
          )
  
      compressed = jac.encode(messages, cdf_fun, cdf_state)
      decompressed = jac.decode(
          compressed,
          jnp.array([message_len] * batch_size),
          inverse_cdf_fun,
          cdf_state,
      )
  
      assert (decompressed == messages).all()
  
file_path: NeuralCompression/projects/jax_entropy_coders/test_entropy_coders.py
imports: import craystack
  import jax.numpy
  import jax_arithmetic_coder
  import numpy
  import pytest
  import scipy
label_desc: Use image node_id 7 for calling a global function with example usage: test_arithmetic_coder_adaptive(shape, alphabet_size) and returns: cdf,  cdf,  freqs, cdf,  cdf,  symbol,  freqs