code: def test_rans_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, message_len, interleave_level = shape
      seed = 7 * batch_size
      rng = np.random.default_rng(seed)
  
      messages = jnp.array(
          rng.integers(
              low=0, high=alphabet_size, size=shape, dtype=np.int64
          )
      )
  
      cdf_state = (
          jnp.array(
              rng.integers(
                  low=1, high=42, size=(batch_size, alphabet_size)
              )
          ),
      )
  
      def cdf_fun(symbols, cdf_state):
          freqs = cdf_state[0]
          cdf = freqs_to_cdf(freqs)
  
          for symbol in symbols:
              freqs.at[symbol].add(1)
  
          return cdf[symbols], cdf[symbols + 1], (freqs,)
  
      def inverse_cdf_fun(cdf_value, cdf_state):
          freqs = cdf_state[0]
          cdf = freqs_to_cdf(freqs)
          symbols = (
              jnp.argmin(jnp.expand_dims(cdf_value, -1) >= cdf, axis=-1)
              - 1
          )
  
          for symbol in jnp.flip(symbols):
              freqs.at[symbol].add(-1)
  
          return cdf[symbols], cdf[symbols + 1], symbols, (freqs,)
  
      codec = craystack.default_rans_codec(
          cdf_fun, inverse_cdf_fun, cdf_state
      )
  
      compressed, cdf_state = craystack.encode(messages, codec)
      codec = craystack.default_rans_codec(
          cdf_fun, inverse_cdf_fun, cdf_state
      )
      decompressed = craystack.decode(
          compressed, message_len, messages.shape[2:], codec
      )[0]
  
      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 11 for calling a global function with example usage: test_rans_coder_adaptive(shape, alphabet_size) and returns: cdf,  cdf,  freqs, cdf,  cdf,  symbols,  freqs