code: def test_distinct_interleaved_freqs(
      batch_size, message_len, interleave_levels, alphabet_size
  ):
      # Tests if codecs can handle different frequency counts for
      # each interleaved level.
      rng = np.random.default_rng(123)
      freqs = jnp.array(
          rng.integers(
              low=1,
              high=10,
              size=(batch_size, interleave_levels, alphabet_size),
          )
      )
      codec = craystack.fixed_array_cdf_codec(freqs_to_cdf(freqs))
      messages = jnp.array(
          rng.integers(
              low=0,
              high=alphabet_size,
              size=(batch_size, message_len, interleave_levels),
              dtype=np.int64,
          )
      )
      decompressed = craystack.decode(
          craystack.encode(messages, codec)[0],
          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 12 for calling a global function with example usage: test_distinct_interleaved_freqs(batch_size, message_len, interleave_levels, alphabet_size) without return types