code: def test_bitsback_ans_codec_identity(
      batch_size,
      message_len,
      interleave_levels,
      obs_alphabet_size,
      latent_alphabet_size,
      message_dtype,
  ):
      # Test if compression is perfectly lossless
      # Uses an identical CDF for all interleave levels,
      # but is a function of symbols and latents
      obs_shape = latent_shape = (interleave_levels,)
  
      rng = np.random.default_rng(123)
      cdf_latent_prior = freqs_to_cdf(  # p(z)
          jnp.array(
              rng.integers(
                  low=1,
                  high=10,
                  size=(*latent_shape, latent_alphabet_size),
              )
          )
      )
      cdf_latent_posterior = freqs_to_cdf(  # q(z|x)
          jnp.array(
              rng.integers(
                  low=1,
                  high=10,
                  size=(obs_alphabet_size, latent_alphabet_size),
              )
          )
      )
      cdf_obs = freqs_to_cdf(  # p(x|z)
          jnp.array(
              rng.integers(
                  low=1,
                  high=10,
                  size=(latent_alphabet_size, obs_alphabet_size),
              )
          )
      )
  
      messages = jnp.array(
          rng.integers(
              low=0,
              high=obs_alphabet_size,
              size=(batch_size, message_len, *obs_shape),
              dtype=message_dtype,
          )
      )
  
      latent_prior_codec = craystack.fixed_array_cdf_codec(
          cdf_latent_prior, message_dtype=message_dtype
      )
  
      def latent_posterior_codec_maker(symbols):
          return craystack.fixed_array_cdf_codec(
              cdf_latent_posterior[symbols],
              allow_empty_pops=True,
              message_dtype=message_dtype,
          )
  
      def obs_codec_maker(latents):
          return craystack.fixed_array_cdf_codec(
              cdf_obs[latents], message_dtype=message_dtype
          )
  
      codec = craystack.bitsback_ans_codec(
          latent_prior_codec,
          latent_posterior_codec_maker,
          obs_codec_maker,
          latent_shape,
          message_dtype,
      )
  
      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 13 for calling a global function with example usage: test_bitsback_ans_codec_identity(batch_size, message_len, interleave_levels, obs_alphabet_size, latent_alphabet_size, message_dtype) and returns: craystack, craystack