code: def test_arithmetic_coder_skewed(
      base_mass, batch_size, alphabet_size
  ):
      # Tests the AC codec with a highly skewed distribution, i.e, Most of
      # the mass is on the first symbol, e.g. freqs = [999, 1, 1, ...]
      # Uses a typical sequence, i.e. empirical count perfectly matches the pdf.
      # This means the empirical entropy is equal to the true entropy
      freqs = jnp.append(
          base_mass - alphabet_size + 1, jnp.ones(alphabet_size - 1)
      ).astype(int)
      messages = jnp.tile(
          jnp.repeat(jnp.arange(alphabet_size), freqs), (batch_size, 1)
      )
      cdf_state = (jnp.tile(freqs_to_cdf(freqs), (batch_size, 1)),)
  
      message_len = messages.shape[1]
      entropy = scipy.stats.entropy(freqs, base=2)
      predicted_message_size = int(np.ceil(entropy / 8 * message_len))
  
      def cdf_fun(symbol, cdf_state):
          cdf = cdf_state[0]
          return cdf[symbol], cdf[symbol + 1], cdf_state
  
      def inverse_cdf_fun(value, cdf_state):
          cdf = cdf_state[0]
          symbol = jnp.argmin(value >= cdf) - 1
          return cdf[symbol], cdf[symbol + 1], symbol, cdf_state
  
      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()
      assert len(compressed[0]) == predicted_message_size
  
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 8 for calling a global function with example usage: test_arithmetic_coder_skewed(base_mass, batch_size, alphabet_size) and returns: cdf,  cdf,  cdf_state, cdf,  cdf,  symbol,  cdf_state