code: def test_arithmetic_coder_identity(shape, alphabet_size):
      # Uses the true source distribution to compress.
      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
          )
      )
      cdfs = jnp.tile(
          calculate_sample_cdf(alphabet_size, rng), (batch_size, 1)
      )
      cdf_state = (jnp.array(cdfs),)
  
      def cdf_fun(symbol, cdf_state):
          return (
              cdf_state[0][symbol],
              cdf_state[0][symbol + 1],
              cdf_state,
          )
  
      def inverse_cdf_fun(value, cdf_state):
          symbol = jnp.argmin(value >= cdf_state[0]) - 1
          return (
              cdf_state[0][symbol],
              cdf_state[0][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()
  
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 5 for calling a global function with example usage: test_arithmetic_coder_identity(shape, alphabet_size) and returns: cdf_state,  cdf_state,  cdf_state, cdf_state,  cdf_state,  symbol,  cdf_state