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: Compresses a floating-point tensor.
  
  Compresses the tensor to bit strings. `bottleneck` is first quantized
  as in `quantize()`, and then compressed using the probability tables in
  `self.cdf` (derived from `self.prior`). The quantized tensor can later be
  recovered by calling `decompress()`.
  
  The innermost `self.coding_rank` dimensions are treated as one coding unit,
  i.e. are compressed into one string each. Any additional dimensions to the
  left are treated as batch dimensions.
  
  Args:
    bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
      least `self.coding_rank` dimensions, and the innermost dimensions must
      be broadcastable to `self.prior_shape`.
  
  Returns:
    A `tf.Tensor` having the same shape as `bottleneck` without the
    `self.coding_rank` innermost dimensions, containing a string for each
    coding unit.
docstr_tok: ["Compresses","a","floating-point","tensor",".","Compresses","the","tensor","to","bit","strings",".","`","bottleneck","`","is","first","quantized","as","in","`","quantize","(",")","`",",","and","then","compressed","using","the","probability","tables","in","`","self.cdf","`","(","derived","from","`","self.prior","`",")",".","The","quantized","tensor","can","later","be","recovered","by","calling","`","decompress","(",")","`",".","The","innermost","`","self.coding_rank","`","dimensions","are","treated","as","one","coding","unit",",","i.e",".","are","compressed","into","one","string","each",".","Any","additional","dimensions","to","the","left","are","treated","as","batch","dimensions",".","Args",":","bottleneck",":","`","tf.Tensor","`","containing","the","data","to","be","compressed",".","Must","have","at","least","`","self.coding_rank","`","dimensions",",","and","the","innermost","dimensions","must","be","broadcastable","to","`","self.prior_shape","`",".","Returns",":","A","`","tf.Tensor","`","having","the","same","shape","as","`","bottleneck","`","without","the","`","self.coding_rank","`","innermost","dimensions",",","containing","a","string","for","each","coding","unit","."]
code: def compress(self, bottleneck):
      """Compresses a floating-point tensor.
  
      Compresses the tensor to bit strings. `bottleneck` is first quantized
      as in `quantize()`, and then compressed using the probability tables in
      `self.cdf` (derived from `self.prior`). The quantized tensor can later be
      recovered by calling `decompress()`.
  
      The innermost `self.coding_rank` dimensions are treated as one coding unit,
      i.e. are compressed into one string each. Any additional dimensions to the
      left are treated as batch dimensions.
  
      Args:
        bottleneck: `tf.Tensor` containing the data to be compressed. Must have at
          least `self.coding_rank` dimensions, and the innermost dimensions must
          be broadcastable to `self.prior_shape`.
  
      Returns:
        A `tf.Tensor` having the same shape as `bottleneck` without the
        `self.coding_rank` innermost dimensions, containing a string for each
        coding unit.
      """
      bottleneck = tf.convert_to_tensor(
          bottleneck, dtype=self.bottleneck_dtype
      )
      input_shape = tf.shape(bottleneck)
      all_but_last_n_elems = lambda t, n: t[:-n] if n else t
      batch_shape = all_but_last_n_elems(input_shape, self.coding_rank)
      iid_shape = all_but_last_n_elems(
          input_shape, self.prior_shape.rank
      )
      offset = self.quantization_offset
      if offset is not None:
          bottleneck -= offset
      symbols = tf.cast(tf.round(bottleneck), tf.int32)
      symbols = tf.reshape(symbols, tf.concat([iid_shape, [-1]], 0))
      symbols -= self.cdf_offset
      handle = gen_ops.create_range_encoder(batch_shape, self.cdf)
      handle = gen_ops.entropy_encode_channel(handle, symbols)
      return gen_ops.entropy_encode_finalize(handle)
  
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 8 for calling the ContinuousBatchedEntropyModel obj's underlying member method code with example usage: obj.compress(bottleneck) and returns: gen_ops