code: def test_gaussian_ub(
      data_min: Tuple[float, float] = (-2.0, -3.0),
      data_max: Tuple[float, float] = (2.0, 3.0),
      bin_prec: int = 10,
      coding_prec: int = 16,
      batch_size: int = 5,
  ):
      """
      Test the codec for the diagonal Gaussian with uniform bins.
  
      Based on `craystack.codecs_test.test_gaussian_ub`
      (https://github.com/j-towns/craystack).
      """
      mean = rng.randn(batch_size, 2) / 10
      stdd = np.exp(rng.random((batch_size, 2)) / 2)
  
      data = np.array(
          [rng.choice(1 << bin_prec, 2) for _ in range(batch_size)]
      )
      check_codec(
          (batch_size, 2),
          diag_gaussian_unif_bins(
              mean,
              stdd,
              np.array(data_min),
              np.array(data_max),
              coding_prec,
              bin_prec,
          ),
          data,
      )
  
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 1 for calling a global function with example usage: test_gaussian_ub(data_min, data_max, bin_prec, coding_prec, batch_size) without return types