code: def test_partial_code_decode(shape1, shape2, alphabet_size):
      # here we test encoding messages into a buffer that already has information
      # i.e., multiple sequential encodes
      tail_capacity = 50000
      batch_size, message_len1, interleave_level1 = shape1
      _, message_len2, interleave_level2 = shape2
      rng = np.random.default_rng(123)
      freqs = jnp.ones(alphabet_size) * message_len1 // alphabet_size
      messages1 = jnp.array(
          rng.integers(
              low=0, high=alphabet_size, size=shape1, dtype=np.uint8
          )
      )
      messages2 = jnp.array(
          rng.integers(
              low=0, high=alphabet_size, size=shape2, dtype=np.uint8
          )
      )
  
      codec1 = craystack.fixed_array_cdf_codec(
          jnp.tile(
              freqs_to_cdf(freqs), (batch_size, interleave_level1, 1)
          ).astype(jnp.uint32)
      )
  
      codec2 = craystack.fixed_array_cdf_codec(
          jnp.tile(
              freqs_to_cdf(freqs), (batch_size, interleave_level2, 1)
          ).astype(jnp.uint32)
      )
  
      compressed = craystack.encode(
          messages1, codec1, tail_capacity=tail_capacity
      )[0]
      compressed = craystack.encode(
          messages2,
          codec2,
          start_buffers=compressed,
          tail_capacity=tail_capacity,
      )[0]
      decompressed2, compressed, _ = craystack.decode(
          compressed,
          message_len2,
          messages2.shape[2:],
          codec2,
          tail_capacity=tail_capacity,
      )
      decompressed1, _, _ = craystack.decode(
          compressed,
          message_len1,
          messages1.shape[2:],
          codec1,
          tail_capacity=tail_capacity,
      )
  
      assert (decompressed1 == messages1).all()
      assert (decompressed2 == messages2).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 10 for calling a global function with example usage: test_partial_code_decode(shape1, shape2, alphabet_size) without return types