code: def diag_gaussian_unif_bins(
      mean: np.ndarray,
      stdd: np.ndarray,
      data_min: Union[float, np.ndarray],
      data_max: Union[float, np.ndarray],
      coding_prec: int,
      bin_prec: int,
  ) -> cs.Codec:
      """
      Codec for data from a diagonal Gaussian with uniform bins.
  
      Based on `craystack.codecs.DiagGaussian_UnifBins`
      (https://github.com/j-towns/craystack)
      to extend first/last bins until minus/plus infinity.
  
      Args:
          mean: Mean of the Gaussian.
          stdd: Standard deviation of the Gaussian.
          data_min: Minimum of the data.
          data_max: Maximum of the data.
          coding_prec: Precision to use for coding.
          bin_prec: The precision to use for the uniform bins.
      Returns:
          The craystack codec for encoding Gaussian data
          discretized using uniform bins.
      """
      n_bins = 1 << bin_prec
      bin_min = data_min - 1 / (n_bins - 1)
      bin_max = data_max + 1 / (n_bins - 1)
  
      bins = np.linspace(bin_min, bin_max, n_bins + 1)
      bins = np.broadcast_to(
          np.moveaxis(bins, 0, -1), mean.shape + (n_bins + 1,)
      )
      cdfs = codecs.norm.cdf(
          bins, mean[..., np.newaxis], stdd[..., np.newaxis]
      )
      cdfs[..., 0] = 0
      cdfs[..., -1] = 1
      pmfs = cdfs[..., 1:] - cdfs[..., :-1]
      buckets = codecs._cumulative_buckets_from_probs(pmfs, coding_prec)
      enc_statfun = codecs._cdf_to_enc_statfun(
          codecs._cdf_from_cumulative_buckets(buckets)
      )
      dec_statfun = codecs._ppf_from_cumulative_buckets(buckets)
  
      return codecs.NonUniform(enc_statfun, dec_statfun, coding_prec)
  
imports: from abc import ABC, abstractmethod
  from typing import List, NamedTuple, Optional, Tuple, Union, overload
  import craystack
  import numpy
  import torch
  from craystack import codecs
  from craystack.bb_ans import BBANS
  from improved_diffusion.gaussian_diffusion import GaussianDiffusion
label_desc: Use image node_id 1 for calling a global function with example usage: diag_gaussian_unif_bins(mean, stdd, data_min, data_max, coding_prec, bin_prec) and returns: codecs