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: Perturbs a tensor with additive uniform noise and estimates bitcost.
  
  Args:
    bottleneck: `tf.Tensor` containing a non-perturbed bottleneck. Must have
      at least `self.coding_rank` dimensions.
    training: Boolean. If `False`, computes the bitcost using discretized
     uniform noise. If `True`, estimates the differential entropy with uniform
     noise.
  
  Returns:
    A tuple
    (bottleneck_perturbed, bits)
    where `bottleneck_perturbed` is `bottleneck` perturbed with nosie
    and `bits` is the bitcost of transmitting such a sample having the same
    shape as `bottleneck` without the `self.coding_rank` innermost dimensions.
docstr_tok: ["Perturbs","a","tensor","with","additive","uniform","noise","and","estimates","bitcost",".","Args",":","bottleneck",":","`","tf.Tensor","`","containing","a","non-perturbed","bottleneck",".","Must","have","at","least","`","self.coding_rank","`","dimensions",".","training",":","Boolean",".","If","`","False","`",",","computes","the","bitcost","using","discretized","uniform","noise",".","If","`","True","`",",","estimates","the","differential","entropy","with","uniform","noise",".","Returns",":","A","tuple","(","bottleneck_perturbed",",","bits",")","where","`","bottleneck_perturbed","`","is","`","bottleneck","`","perturbed","with","nosie","and","`","bits","`","is","the","bitcost","of","transmitting","such","a","sample","having","the","same","shape","as","`","bottleneck","`","without","the","`","self.coding_rank","`","innermost","dimensions","."]
code: def __call__(self, bottleneck, training=True):
      """Perturbs a tensor with additive uniform noise and estimates bitcost.
  
      Args:
        bottleneck: `tf.Tensor` containing a non-perturbed bottleneck. Must have
          at least `self.coding_rank` dimensions.
        training: Boolean. If `False`, computes the bitcost using discretized
         uniform noise. If `True`, estimates the differential entropy with uniform
         noise.
  
      Returns:
        A tuple
        (bottleneck_perturbed, bits)
        where `bottleneck_perturbed` is `bottleneck` perturbed with nosie
        and `bits` is the bitcost of transmitting such a sample having the same
        shape as `bottleneck` without the `self.coding_rank` innermost dimensions.
      """
      bottleneck = tf.convert_to_tensor(
          bottleneck, dtype=self.bottleneck_dtype
      )
      log_prob_fn = functools.partial(self._log_prob, self.prior)
      if training:
          log_probs, bottleneck_perturbed = math_ops.perturb_and_apply(
              log_prob_fn,
              bottleneck,
              expected_grads=self._expected_grads,
          )
      else:
          # Here we compute `H(round(bottleneck - noise) | noise )`.
          input_shape = tf.shape(bottleneck)
          input_rank = tf.shape(input_shape)[0]
          _, coding_shape = tf.split(
              input_shape,
              [input_rank - self.coding_rank, self.coding_rank],
          )
          broadcast_shape = coding_shape[
              : self.coding_rank - self.prior_shape.rank
          ]
          _, offset = self._compute_indexes_and_offset(broadcast_shape)
          symbols = tf.round(bottleneck - offset)
          bottleneck_perturbed = symbols + offset
          log_probs = log_prob_fn(bottleneck_perturbed)
  
      axes = tuple(range(-self.coding_rank, 0))
      bits = tf.reduce_sum(log_probs, axis=axes) / (
          -tf.math.log(tf.constant(2.0, dtype=log_probs.dtype))
      )
      return bottleneck_perturbed, bits
  
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 5 for calling the UniversalBatchedEntropyModel obj's underlying member method code with example usage: obj.__call__(bottleneck, training) and returns: bottleneck_perturbed,  bits