class_docstr: Batched entropy model for continuous random variables.
  
  This entropy model handles quantization of a bottleneck tensor and helps with
  training of the parameters of the probability distribution modeling the
  tensor (a shared "prior" between sender and receiver). It also pre-computes
  integer probability tables, which can then be used to compress and decompress
  bottleneck tensors reliably across different platforms.
  
  A typical workflow looks like this:
  
  - Train a model using an instance of this entropy model as a bottleneck,
    passing the bottleneck tensor through it. With `training=True`, the model
    computes a differentiable upper bound on the number of bits needed to
    compress the bottleneck tensor.
  - For evaluation, get a closer estimate of the number of compressed bits
    using `training=False`.
  - Instantiate an entropy model with `compression=True` (and the same
    parameters as during training), and share the model between a sender and a
    receiver.
  - On the sender side, compute the bottleneck tensor and call `compress()` on
    it. The output is a compressed string representation of the tensor. Transmit
    the string to the receiver, and call `decompress()` there. The output is the
    quantized bottleneck tensor. Continue processing the tensor on the receiving
    side.
  
  This class assumes that all scalar elements of the encoded tensor are
  statistically independent, and that the parameters of their scalar
  distributions do not depend on data. The innermost dimensions of the
  bottleneck tensor must be broadcastable to the batch shape of `prior`. Any
  dimensions to the left of the batch shape are assumed to be i.i.d., i.e. the
  likelihoods are broadcast to the bottleneck tensor accordingly.
  
  A more detailed description (and motivation) of this way of performing
  quantization and range coding can be found in the following paper. Please cite
  the paper when using this code for derivative work.
  
  > "End-to-end Optimized Image Compression"<br />
  > J. Ballé, V. Laparra, E.P. Simoncelli<br />
  > https://openreview.net/forum?id=rJxdQ3jeg
  
  Entropy models which contain range coding tables (i.e. with
  `compression=True`) can be instantiated in three ways:
  
  - By providing a continuous "prior" distribution object. The range coding
    tables are then derived from that continuous distribution.
  - From a config as returned by `get_config`, followed by a call to
    `set_weights`. This implements the Keras serialization protocol. In this
    case, the initializer creates empty state variables for the range coding
    tables, which are then filled by `set_weights`. As a consequence, this
    method requires `stateless=False`.
  - In a more low-level way, by directly providing the range coding tables to
    `__init__`, for use cases where the Keras protocol can't be used (e.g., when
    the entropy model must not create variables).
  
  The `quantization_offset` and `offset_heuristic` arguments control whether
  quantization is performed with respect to integer values, or potentially
  non-integer offsets (i.e., `y = tf.round(x - o) + o`). There are three modes
  of operation:
  
  - If `quantization_offset` is provided manually (not `None`), these values are
    used and `offset_heuristic` is ineffective.
  - Otherwise, if `offset_heuristic and compression`, the offsets are computed
    once on initialization, and then fixed. If the entropy model is serialized,
    they are preserved.
  - Otherwise, if `offset_heuristic and not compression`, the offsets are
    recomputed every time quantization is performed. Note this may be
    computationally expensive when the prior does not have a mode that is
    computable in closed form (e.g. for `NoisyDeepFactorized`).
  
  This offset heuristic is discussed in Section III.A of:
  > "Nonlinear Transform Coding"<br />
  > J. Ballé, P.A. Chou, D. Minnen, S. Singh, N. Johnston, E. Agustsson,
  > S.J. Hwang, G. Toderici<br />
  > https://doi.org/10.1109/JSTSP.2020.3034501
class_docstr_tok: ["Batched","entropy","model","for","continuous","random","variables",".","This","entropy","model","handles","quantization","of","a","bottleneck","tensor","and","helps","with","training","of","the","parameters","of","the","probability","distribution","modeling","the","tensor","(","a","shared","``","prior","''","between","sender","and","receiver",")",".","It","also","pre-computes","integer","probability","tables",",","which","can","then","be","used","to","compress","and","decompress","bottleneck","tensors","reliably","across","different","platforms",".","A","typical","workflow","looks","like","this",":","-","Train","a","model","using","an","instance","of","this","entropy","model","as","a","bottleneck",",","passing","the","bottleneck","tensor","through","it",".","With","`","training=True","`",",","the","model","computes","a","differentiable","upper","bound","on","the","number","of","bits","needed","to","compress","the","bottleneck","tensor",".","-","For","evaluation",",","get","a","closer","estimate","of","the","number","of","compressed","bits","using","`","training=False","`",".","-","Instantiate","an","entropy","model","with","`","compression=True","`","(","and","the","same","parameters","as","during","training",")",",","and","share","the","model","between","a","sender","and","a","receiver",".","-","On","the","sender","side",",","compute","the","bottleneck","tensor","and","call","`","compress","(",")","`","on","it",".","The","output","is","a","compressed","string","representation","of","the","tensor",".","Transmit","the","string","to","the","receiver",",","and","call","`","decompress","(",")","`","there",".","The","output","is","the","quantized","bottleneck","tensor",".","Continue","processing","the","tensor","on","the","receiving","side",".","This","class","assumes","that","all","scalar","elements","of","the","encoded","tensor","are","statistically","independent",",","and","that","the","parameters","of","their","scalar","distributions","do","not","depend","on","data",".","The","innermost","dimensions","of","the","bottleneck","tensor","must","be","broadcastable","to","the","batch","shape","of","`","prior","`",".","Any","dimensions","to","the","left","of","the","batch","shape","are","assumed","to","be","i.i.d.",",","i.e",".","the","likelihoods","are","broadcast","to","the","bottleneck","tensor","accordingly",".","A","more","detailed","description","(","and","motivation",")","of","this","way","of","performing","quantization","and","range","coding","can","be","found","in","the","following","paper",".","Please","cite","the","paper","when","using","this","code","for","derivative","work",".",">","``","End-to-end","Optimized","Image","Compression","''","<","br","\/",">",">","J.","Ball\u00e9",",","V.","Laparra",",","E.P",".","Simoncelli","<","br","\/",">",">","https",":","\/\/openreview.net\/forum","?","id=rJxdQ3jeg","Entropy","models","which","contain","range","coding","tables","(","i.e",".","with","`","compression=True","`",")","can","be","instantiated","in","three","ways",":","-","By","providing","a","continuous","``","prior","''","distribution","object",".","The","range","coding","tables","are","then","derived","from","that","continuous","distribution",".","-","From","a","config","as","returned","by","`","get_config","`",",","followed","by","a","call","to","`","set_weights","`",".","This","implements","the","Keras","serialization","protocol",".","In","this","case",",","the","initializer","creates","empty","state","variables","for","the","range","coding","tables",",","which","are","then","filled","by","`","set_weights","`",".","As","a","consequence",",","this","method","requires","`","stateless=False","`",".","-","In","a","more","low-level","way",",","by","directly","providing","the","range","coding","tables","to","`","__init__","`",",","for","use","cases","where","the","Keras","protocol","ca","n't","be","used","(","e.g.",",","when","the","entropy","model","must","not","create","variables",")",".","The","`","quantization_offset","`","and","`","offset_heuristic","`","arguments","control","whether","quantization","is","performed","with","respect","to","integer","values",",","or","potentially","non-integer","offsets","(","i.e.",",","`","y","=","tf.round","(","x","-","o",")","+","o","`",")",".","There","are","three","modes","of","operation",":","-","If","`","quantization_offset","`","is","provided","manually","(","not","`","None","`",")",",","these","values","are","used","and","`","offset_heuristic","`","is","ineffective",".","-","Otherwise",",","if","`","offset_heuristic","and","compression","`",",","the","offsets","are","computed","once","on","initialization",",","and","then","fixed",".","If","the","entropy","model","is","serialized",",","they","are","preserved",".","-","Otherwise",",","if","`","offset_heuristic","and","not","compression","`",",","the","offsets","are","recomputed","every","time","quantization","is","performed",".","Note","this","may","be","computationally","expensive","when","the","prior","does","not","have","a","mode","that","is","computable","in","closed","form","(","e.g",".","for","`","NoisyDeepFactorized","`",")",".","This","offset","heuristic","is","discussed","in","Section","III.A","of",":",">","``","Nonlinear","Transform","Coding","''","<","br","\/",">",">","J.","Ball\u00e9",",","P.A",".","Chou",",","D.","Minnen",",","S.","Singh",",","N.","Johnston",",","E.","Agustsson",",",">","S.J",".","Hwang",",","G.","Toderici","<","br","\/",">",">","https",":","\/\/doi.org\/10.1109\/JSTSP.2020.3034501"]
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 in the __call__ method are summed over each coding unit.
    compression: Boolean. If set to `True`, the range coding tables used by
      `compress()` and `decompress()` will be built on instantiation. If set
      to `False`, these two methods will not be accessible.
    stateless: Boolean. If `False`, range coding tables are created as
      `Variable`s. This allows the entropy model to be serialized using the
      `SavedModel` protocol, so that both the encoder and the decoder use
      identical tables when loading the stored model. If `True`, creates range
      coding tables as `Tensor`s. This makes the entropy model stateless and
      allows it to be constructed within a `tf.function` body, for when the
      range coding tables are provided manually. If `compression=False`, then
      `stateless=True` is implied and the provided value is ignored.
    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`.
    prior_shape: Batch shape of the prior (dimensions which are not assumed
      i.i.d.). Must be provided if `prior` is omitted.
    cdf: `tf.Tensor` or `None`. If provided, is used for range coding rather
      than tables built from the prior.
    cdf_offset: `tf.Tensor` or `None`. Must be provided along with `cdf`.
    cdf_shapes: Shapes of `cdf` and `cdf_offset`. If provided, empty range
      coding tables are created, which can then be restored using
      `set_weights`. Requires `compression=True` and `stateless=False`.
    offset_heuristic: Boolean. Whether to quantize to non-integer offsets
      heuristically determined from mode/median of prior. Set this to `False`
      if you are using soft quantization during training.
    quantization_offset: `tf.Tensor` or `None`. The quantization offsets to
      use. If provided (not `None`), then `offset_heuristic` is ineffective.
    decode_sanity_check: Boolean. If `True`, an raises an error if the binary
      strings passed into `decompress` are not completely decoded.
    laplace_tail_mass: Float, or a float-valued tf.Tensor. If positive,
      will augment the prior with a `NoisyLaplace` mixture component for
      training stability. (experimental)
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","in","the","__call__","method","are","summed","over","each","coding","unit",".","compression",":","Boolean",".","If","set","to","`","True","`",",","the","range","coding","tables","used","by","`","compress","(",")","`","and","`","decompress","(",")","`","will","be","built","on","instantiation",".","If","set","to","`","False","`",",","these","two","methods","will","not","be","accessible",".","stateless",":","Boolean",".","If","`","False","`",",","range","coding","tables","are","created","as","`","Variable","`","s",".","This","allows","the","entropy","model","to","be","serialized","using","the","`","SavedModel","`","protocol",",","so","that","both","the","encoder","and","the","decoder","use","identical","tables","when","loading","the","stored","model",".","If","`","True","`",",","creates","range","coding","tables","as","`","Tensor","`","s",".","This","makes","the","entropy","model","stateless","and","allows","it","to","be","constructed","within","a","`","tf.function","`","body",",","for","when","the","range","coding","tables","are","provided","manually",".","If","`","compression=False","`",",","then","`","stateless=True","`","is","implied","and","the","provided","value","is","ignored",".","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","`",".","prior_shape",":","Batch","shape","of","the","prior","(","dimensions","which","are","not","assumed","i.i.d.",")",".","Must","be","provided","if","`","prior","`","is","omitted",".","cdf",":","`","tf.Tensor","`","or","`","None","`",".","If","provided",",","is","used","for","range","coding","rather","than","tables","built","from","the","prior",".","cdf_offset",":","`","tf.Tensor","`","or","`","None","`",".","Must","be","provided","along","with","`","cdf","`",".","cdf_shapes",":","Shapes","of","`","cdf","`","and","`","cdf_offset","`",".","If","provided",",","empty","range","coding","tables","are","created",",","which","can","then","be","restored","using","`","set_weights","`",".","Requires","`","compression=True","`","and","`","stateless=False","`",".","offset_heuristic",":","Boolean",".","Whether","to","quantize","to","non-integer","offsets","heuristically","determined","from","mode\/median","of","prior",".","Set","this","to","`","False","`","if","you","are","using","soft","quantization","during","training",".","quantization_offset",":","`","tf.Tensor","`","or","`","None","`",".","The","quantization","offsets","to","use",".","If","provided","(","not","`","None","`",")",",","then","`","offset_heuristic","`","is","ineffective",".","decode_sanity_check",":","Boolean",".","If","`","True","`",",","an","raises","an","error","if","the","binary","strings","passed","into","`","decompress","`","are","not","completely","decoded",".","laplace_tail_mass",":","Float",",","or","a","float-valued","tf.Tensor",".","If","positive",",","will","augment","the","prior","with","a","`","NoisyLaplace","`","mixture","component","for","training","stability",".","(","experimental",")"]
code: def __init__(
      self,
      prior=None,
      coding_rank=None,
      compression=False,
      stateless=False,
      expected_grads=False,
      tail_mass=2**-8,
      range_coder_precision=12,
      bottleneck_dtype=None,
      prior_shape=None,
      cdf=None,
      cdf_offset=None,
      cdf_shapes=None,
      offset_heuristic=True,
      quantization_offset=None,
      decode_sanity_check=True,
      laplace_tail_mass=0,
  ):
      """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 in the __call__ method are summed over each coding unit.
        compression: Boolean. If set to `True`, the range coding tables used by
          `compress()` and `decompress()` will be built on instantiation. If set
          to `False`, these two methods will not be accessible.
        stateless: Boolean. If `False`, range coding tables are created as
          `Variable`s. This allows the entropy model to be serialized using the
          `SavedModel` protocol, so that both the encoder and the decoder use
          identical tables when loading the stored model. If `True`, creates range
          coding tables as `Tensor`s. This makes the entropy model stateless and
          allows it to be constructed within a `tf.function` body, for when the
          range coding tables are provided manually. If `compression=False`, then
          `stateless=True` is implied and the provided value is ignored.
        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`.
        prior_shape: Batch shape of the prior (dimensions which are not assumed
          i.i.d.). Must be provided if `prior` is omitted.
        cdf: `tf.Tensor` or `None`. If provided, is used for range coding rather
          than tables built from the prior.
        cdf_offset: `tf.Tensor` or `None`. Must be provided along with `cdf`.
        cdf_shapes: Shapes of `cdf` and `cdf_offset`. If provided, empty range
          coding tables are created, which can then be restored using
          `set_weights`. Requires `compression=True` and `stateless=False`.
        offset_heuristic: Boolean. Whether to quantize to non-integer offsets
          heuristically determined from mode/median of prior. Set this to `False`
          if you are using soft quantization during training.
        quantization_offset: `tf.Tensor` or `None`. The quantization offsets to
          use. If provided (not `None`), then `offset_heuristic` is ineffective.
        decode_sanity_check: Boolean. If `True`, an raises an error if the binary
          strings passed into `decompress` are not completely decoded.
        laplace_tail_mass: Float, or a float-valued tf.Tensor. If positive,
          will augment the prior with a `NoisyLaplace` mixture component for
          training stability. (experimental)
      """
      if (prior is None) == (prior_shape is None):
          raise ValueError(
              "Either `prior` or `prior_shape` must be provided."
          )
      if (prior is None) + (cdf_shapes is None) + (cdf is None) != 2:
          raise ValueError(
              "Must provide exactly one of `prior`, `cdf`, or `cdf_shapes`."
          )
      if not compression and not (
          cdf is None and cdf_offset is None and cdf_shapes is None
      ):
          raise ValueError(
              "CDFs can't be provided with `compression=False`"
          )
      if prior is not None and 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._offset_heuristic = bool(offset_heuristic)
      self._prior_shape = tf.TensorShape(
          prior_shape if prior is None else prior.batch_shape
      )
      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 cdf_shapes is not None:
              # `cdf_shapes` being set indicates that we are using the `SavedModel`
              # protocol, which can only provide JSON datatypes. So create a
              # placeholder value depending on whether `quantization_offset` was
              # `None` or not. For this purpose, we expect a Boolean (when in all
              # other cases, we expect either `None` or a tensor).
              assert isinstance(quantization_offset, bool)
              assert self.compression
              if quantization_offset:
                  quantization_offset = tf.zeros(
                      self.prior_shape_tensor
                  )
              else:
                  quantization_offset = None
          elif quantization_offset is not None:
              # If quantization offset is passed in manually, use it.
              pass
          elif self.offset_heuristic and self.compression:
              # For compression, we need to fix the offset value, so compute it here.
              if self._prior is None:
                  raise ValueError(
                      "To use the offset heuristic, a `prior` needs to be provided."
                  )
              quantization_offset = helpers.quantization_offset(
                  self.prior
              )
              # Optimization: if the quantization offset is zero, we don't need to
              # subtract/add it when quantizing, and we don't need to serialize its
              # value. Note that this code will only work in eager mode.
              if tf.executing_eagerly() and tf.reduce_all(
                  tf.equal(quantization_offset, 0.0)
              ):
                  quantization_offset = None
              else:
                  quantization_offset = tf.broadcast_to(
                      quantization_offset, self.prior_shape_tensor
                  )
          else:
              quantization_offset = None
          if quantization_offset is None:
              self._quantization_offset = None
          elif self.compression and not self.stateless:
              quantization_offset = tf.cast(
                  quantization_offset, self.bottleneck_dtype
              )
              self._quantization_offset = tf.Variable(
                  quantization_offset,
                  trainable=False,
                  name="quantization_offset",
              )
          else:
              quantization_offset = tf.cast(
                  quantization_offset, self.bottleneck_dtype
              )
              self._quantization_offset = tf.convert_to_tensor(
                  quantization_offset, name="quantization_offset"
              )
          if self.compression:
              if cdf is None and cdf_shapes is None:
                  cdf, cdf_offset = self._build_tables(
                      self.prior,
                      range_coder_precision,
                      offset=quantization_offset,
                  )
              self._init_compression(cdf, cdf_offset, cdf_shapes)
  
file_path: compression/tensorflow_compression/python/entropy_models/continuous_batched.py
imports: import functools
  from absl import logging
  import tensorflow
  from tensorflow_compression.python.distributions import helpers
  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
  from tensorflow_compression.python.ops import round_ops
label_desc: Use image node_id 1 to create a new ContinuousBatchedEntropyModel object from inherited base classes: continuous_base with example: obj = ContinuousBatchedEntropyModel(prior, coding_rank, compression, stateless, expected_grads, tail_mass, range_coder_precision, bottleneck_dtype, prior_shape, cdf, cdf_offset, cdf_shapes, offset_heuristic, quantization_offset, decode_sanity_check, laplace_tail_mass)