code: def test_codec(
      batch_size: int = 2,
      image_size: int = 3,
      data_prec: int = 8,
      data_min: float = -1,
      data_max: float = 1,
      steps: int = 40,
  ):
      """
      Test the Bits-Back codec for diffusion models.
      """
      data_shape = (batch_size, 3, image_size, image_size)
      diffusion = create_gaussian_diffusion(
          timestep_respacing=f"vlb{steps}"
      )
      discr_data = rng.randint(
          low=0, high=(1 << data_prec) - 1, size=data_shape
      )
      data = data_min + discr_data.astype(float) * (
          data_max - data_min
      ) / ((1 << data_prec) - 1)
      codec_model = CheatingDiffusionModel(
          data, diffusion=diffusion, model=torch.nn.Linear(1, 1)
      )
      codec = BitsBackCodec(
          model=codec_model,
          data_shape=data_shape,
          data_min=data_min,
          data_max=data_max,
          track=True,
      )
      message = deepcopy(codec.message)
      codec.encode(discr_data)
      stats = codec.statistics()
      rate_extra = codec._bytes_to_bpd(codec.extra_bytes)
      _, decoded_data = codec.decode()
  
      assert stats.rate_effective_push is not None
      assert stats.rate_effective_pop is not None
      np.testing.assert_equal(discr_data, decoded_data)
      np.testing.assert_allclose(
          stats.rate_effective_push, stats.rate_effective_pop
      )
  
      assert rans.message_equal(message, codec.message)
      assert codec.flat_message.nbytes == codec.initial_bytes
      assert stats.samples == batch_size
      assert np.isclose(
          stats.rate_effective, stats.rate_effective_push[0]
      )
      assert np.isclose(stats.rate_effective, 0.0)
      assert np.isclose(stats.rate_used, rate_extra)
      assert np.isclose(stats.rate_used, sum(stats.rate_effective_push))
      assert codec.data_count == 0
  
file_path: NeuralCompression/projects/bits_back_diffusion/test.py
imports: from copy import deepcopy
  from typing import Optional, Tuple
  import numpy
  import numpy.random
  import torch
  from bits_back_diffusion.codec import BitsBackCodec, DiffusionModel, GaussianParams, diag_gaussian_unif_bins
  from bits_back_diffusion.script_util import create_gaussian_diffusion
  from craystack import rans
  from craystack.codecs_test import check_codec
label_desc: Use image node_id 2 for calling a global function with example usage: test_codec(batch_size, image_size, data_prec, data_min, data_max, steps) without return types