code: def test_craystack_coder_identity(shape, alphabet_size, skewed):
      # here we only test identity and not message length
      # due to implementation details we don't exactly achieve the
      # entropy-predicted len
      batch_size, message_len, interleave_level = shape
      if skewed:
          freqs = jnp.array(
              generate_skewed_distribution(
                  alphabet_size, message_len, 124
              )
          )
          messages = jnp.tile(
              jnp.repeat(jnp.arange(alphabet_size), freqs)[
                  None, :, None
              ],
              (batch_size, 1, interleave_level),
          )
      else:
          rng = np.random.default_rng(123)
          freqs = jnp.ones(alphabet_size) * message_len // alphabet_size
          messages = jnp.array(
              rng.integers(
                  low=0, high=alphabet_size, size=shape, dtype=np.uint8
              )
          )
  
      codec = craystack.fixed_array_cdf_codec(
          jnp.tile(
              freqs_to_cdf(freqs), (batch_size, interleave_level, 1)
          ).astype(jnp.uint32)
      )
  
      compressed = craystack.encode(messages, codec)[0]
      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 9 for calling a global function with example usage: test_craystack_coder_identity(shape, alphabet_size, skewed) without return types