class_docstr: Batched entropy model model which implements Universal Quantization.
  
  In contrast to the base class, which uses rounding for quantization, here
  "quantization" is performed additive uniform noise, which is implemented with
  Universal Quantization.
  
  This is described in Sec. 3.2. in the paper
  > "Universally Quantized Neural Compression"<br />
  > Eirikur Agustsson & Lucas Theis<br />
  > https://arxiv.org/abs/2006.09952
class_docstr_tok: ["Batched","entropy","model","model","which","implements","Universal","Quantization",".","In","contrast","to","the","base","class",",","which","uses","rounding","for","quantization",",","here","''","quantization","''","is","performed","additive","uniform","noise",",","which","is","implemented","with","Universal","Quantization",".","This","is","described","in","Sec",".","3.2.","in","the","paper",">","``","Universally","Quantized","Neural","Compression","''","<","br","\/",">",">","Eirikur","Agustsson","&","Lucas","Theis","<","br","\/",">",">","https",":","\/\/arxiv.org\/abs\/2006.09952"]
docstr: Initializes the instance.
  
  Args:
    prior: A `tfp.distributions.Distribution` object. A density model fitting
      the marginal distribution of the bottleneck data with additive uniform
      noise, which is shared a priori between the sender and the receiver. For
      best results, the distribution should be flexible enough to have a
      unit-width uniform distribution as a special case, since this is the
      marginal distribution for bottleneck dimensions that are constant. The
      distribution parameters may not depend on data (they must be either
      variables or constants).
    coding_rank: Integer. Number of innermost dimensions considered a coding
      unit. Each coding unit is compressed to its own bit string, and the
      `bits()` method sums over each coding unit.
    compression: Boolean. If set to `True`, the range coding tables used by
      `compress()` and `decompress()` will be built on instantiation. This
      assumes eager mode (throws an error if in graph mode or inside a
      `tf.function` call). If set to `False`, these two methods will not be
      accessible.
    laplace_tail_mass: Float. If positive, will augment the prior with a
      laplace mixture for training stability.
    expected_grads: If True, will use analytical expected gradients during
      backpropagation w.r.t. additive uniform noise.
    tail_mass: Float. Approximate probability mass which is encoded using an
      Elias gamma code embedded into the range coder.
    range_coder_precision: Integer. Precision passed to the range coding op.
    bottleneck_dtype: `tf.dtypes.DType`. Data type of bottleneck tensor.
      Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
    num_noise_levels: Integer. The number of levels used to quantize the
      uniform noise.
    stateless: Boolean. If `True`, creates range coding tables as `Tensor`s
      rather than `Variable`s. This makes the entropy model stateless and
      allows it to be constructed within a `tf.function` body. If
      `compression=False`, then `stateless=True` is implied and the provided
      value is ignored.
    decode_sanity_check: Boolean. If `True`, an raises an error if the binary
      strings passed into `decompress` are not completely decoded.
docstr_tok: ["Initializes","the","instance",".","Args",":","prior",":","A","`","tfp.distributions.Distribution","`","object",".","A","density","model","fitting","the","marginal","distribution","of","the","bottleneck","data","with","additive","uniform","noise",",","which","is","shared","a","priori","between","the","sender","and","the","receiver",".","For","best","results",",","the","distribution","should","be","flexible","enough","to","have","a","unit-width","uniform","distribution","as","a","special","case",",","since","this","is","the","marginal","distribution","for","bottleneck","dimensions","that","are","constant",".","The","distribution","parameters","may","not","depend","on","data","(","they","must","be","either","variables","or","constants",")",".","coding_rank",":","Integer",".","Number","of","innermost","dimensions","considered","a","coding","unit",".","Each","coding","unit","is","compressed","to","its","own","bit","string",",","and","the","`","bits","(",")","`","method","sums","over","each","coding","unit",".","compression",":","Boolean",".","If","set","to","`","True","`",",","the","range","coding","tables","used","by","`","compress","(",")","`","and","`","decompress","(",")","`","will","be","built","on","instantiation",".","This","assumes","eager","mode","(","throws","an","error","if","in","graph","mode","or","inside","a","`","tf.function","`","call",")",".","If","set","to","`","False","`",",","these","two","methods","will","not","be","accessible",".","laplace_tail_mass",":","Float",".","If","positive",",","will","augment","the","prior","with","a","laplace","mixture","for","training","stability",".","expected_grads",":","If","True",",","will","use","analytical","expected","gradients","during","backpropagation","w.r.t",".","additive","uniform","noise",".","tail_mass",":","Float",".","Approximate","probability","mass","which","is","encoded","using","an","Elias","gamma","code","embedded","into","the","range","coder",".","range_coder_precision",":","Integer",".","Precision","passed","to","the","range","coding","op",".","bottleneck_dtype",":","`","tf.dtypes.DType","`",".","Data","type","of","bottleneck","tensor",".","Defaults","to","`","tf.keras.mixed_precision.global_policy","(",")",".compute_dtype","`",".","num_noise_levels",":","Integer",".","The","number","of","levels","used","to","quantize","the","uniform","noise",".","stateless",":","Boolean",".","If","`","True","`",",","creates","range","coding","tables","as","`","Tensor","`","s","rather","than","`","Variable","`","s",".","This","makes","the","entropy","model","stateless","and","allows","it","to","be","constructed","within","a","`","tf.function","`","body",".","If","`","compression=False","`",",","then","`","stateless=True","`","is","implied","and","the","provided","value","is","ignored",".","decode_sanity_check",":","Boolean",".","If","`","True","`",",","an","raises","an","error","if","the","binary","strings","passed","into","`","decompress","`","are","not","completely","decoded","."]
code: def __init__(
      self,
      prior,
      coding_rank,
      compression=False,
      laplace_tail_mass=0.0,
      expected_grads=False,
      tail_mass=2**-8,
      range_coder_precision=12,
      bottleneck_dtype=None,
      num_noise_levels=15,
      stateless=False,
      decode_sanity_check=True,
  ):
      """Initializes the instance.
  
      Args:
        prior: A `tfp.distributions.Distribution` object. A density model fitting
          the marginal distribution of the bottleneck data with additive uniform
          noise, which is shared a priori between the sender and the receiver. For
          best results, the distribution should be flexible enough to have a
          unit-width uniform distribution as a special case, since this is the
          marginal distribution for bottleneck dimensions that are constant. The
          distribution parameters may not depend on data (they must be either
          variables or constants).
        coding_rank: Integer. Number of innermost dimensions considered a coding
          unit. Each coding unit is compressed to its own bit string, and the
          `bits()` method sums over each coding unit.
        compression: Boolean. If set to `True`, the range coding tables used by
          `compress()` and `decompress()` will be built on instantiation. This
          assumes eager mode (throws an error if in graph mode or inside a
          `tf.function` call). If set to `False`, these two methods will not be
          accessible.
        laplace_tail_mass: Float. If positive, will augment the prior with a
          laplace mixture for training stability.
        expected_grads: If True, will use analytical expected gradients during
          backpropagation w.r.t. additive uniform noise.
        tail_mass: Float. Approximate probability mass which is encoded using an
          Elias gamma code embedded into the range coder.
        range_coder_precision: Integer. Precision passed to the range coding op.
        bottleneck_dtype: `tf.dtypes.DType`. Data type of bottleneck tensor.
          Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
        num_noise_levels: Integer. The number of levels used to quantize the
          uniform noise.
        stateless: Boolean. If `True`, creates range coding tables as `Tensor`s
          rather than `Variable`s. This makes the entropy model stateless and
          allows it to be constructed within a `tf.function` body. If
          `compression=False`, then `stateless=True` is implied and the provided
          value is ignored.
        decode_sanity_check: Boolean. If `True`, an raises an error if the binary
          strings passed into `decompress` are not completely decoded.
      """
      if prior.event_shape.rank:
          raise ValueError(
              "`prior` must be a (batch of) scalar distribution(s)."
          )
  
      super().__init__(
          coding_rank=coding_rank,
          compression=compression,
          stateless=stateless,
          expected_grads=expected_grads,
          tail_mass=tail_mass,
          bottleneck_dtype=bottleneck_dtype,
          laplace_tail_mass=laplace_tail_mass,
      )
      self._prior = prior
      self._num_noise_levels = num_noise_levels
      if self.coding_rank < self.prior_shape.rank:
          raise ValueError(
              "`coding_rank` can't be smaller than `prior_shape`."
          )
      self.decode_sanity_check = decode_sanity_check
  
      with self.name_scope:
          if self.compression:
              offset = _range_coding_offsets(
                  self._num_noise_levels,
                  self.prior_shape,
                  self.bottleneck_dtype,
              )
              cdf, cdf_offset = self._build_tables(
                  self.prior, range_coder_precision, offset=offset
              )
              self._init_compression(cdf, cdf_offset, None)
  
imports: import functools
  import tensorflow
  from tensorflow_compression.python.entropy_models import continuous_base
  from tensorflow_compression.python.ops import gen_ops
  from tensorflow_compression.python.ops import math_ops
label_desc: Use image node_id 1 to create a new UniversalBatchedEntropyModel object from inherited base classes: continuous_base with example: obj = UniversalBatchedEntropyModel(prior, coding_rank, compression, laplace_tail_mass, expected_grads, tail_mass, range_coder_precision, bottleneck_dtype, num_noise_levels, stateless, decode_sanity_check)