class_docstr: Indexed 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
  conditionally independent given some other random variable, possibly depending
  on data. All dependencies must be represented by the `indexes` tensor. For
  each bottleneck tensor element, it selects the appropriate scalar
  distribution.
  
  The `indexes` tensor must contain only integer values in a pre-specified range
  (but may have floating-point type for purposes of backpropagation). To make
  the distribution conditional on `n`-dimensional indexes, `index_ranges` must
  be specified as an iterable of `n` integers. `indexes` must have the same
  shape as the bottleneck tensor with an additional channel dimension of length
  `n`. The position of the channel dimension is given by `channel_axis`. The
  index values in the `k`th channel must be in the range `[0, index_ranges[k])`.
  If `index_ranges` has only one element (i.e. `n == 1`), `channel_axis` may be
  `None`. In that case, the additional channel dimension is omitted, and the
  `indexes` tensor must have the same shape as the bottleneck tensor.
  
  The implied distribution for the bottleneck tensor is determined as:
  ```
  prior_fn(**{k: f(indexes) for k, f in parameter_fns.items()})
  ```
  
  A more detailed description (and motivation) of this indexing scheme can be
  found in the following paper. Please cite the paper when using this code for
  derivative work.
  
  > "Integer Networks for Data Compression with Latent-Variable Models"<br />
  > J. Ballé, N. Johnston, D. Minnen<br />
  > https://openreview.net/forum?id=S1zz2i0cY7
  
  Examples:
  
  To make a parameterized zero-mean normal distribution, one could use:
  ```
  tfc.ContinuousIndexedEntropyModel(
      prior_fn=tfc.NoisyNormal,
      index_ranges=(64,),
      parameter_fns=dict(
          loc=lambda _: 0.,
          scale=lambda i: tf.exp(i / 8 - 5),
      ),
      coding_rank=1,
      channel_axis=None,
  )
  ```
  Then, each element of `indexes` in the range `[0, 64)` would indicate that the
  corresponding element in `bottleneck` is normally distributed with zero mean
  and a standard deviation between `exp(-5)` and `exp(2.875)`, inclusive.
  
  To make a parameterized logistic mixture distribution, one could use:
  ```
  tfc.ContinuousIndexedEntropyModel(
      prior_fn=tfc.NoisyLogisticMixture,
      index_ranges=(10, 10, 5),
      parameter_fns=dict(
          loc=lambda i: i[..., 0:2] - 5,
          scale=lambda _: 1,
          weight=lambda i: tf.nn.softmax((i[..., 2:3] - 2) * [-1, 1]),
      ),
      coding_rank=1,
      channel_axis=-1,
  )
  ```
  Then, the last dimension of `indexes` would consist of triples of elements in
  the ranges `[0, 10)`, `[0, 10)`, and `[0, 5)`, respectively. Each triple
  would indicate that the element in `bottleneck` corresponding to the other
  dimensions is distributed with a mixture of two logistic distributions, where
  the components each have one of 10 location parameters between `-5` and `+4`,
  inclusive, unit scale parameters, and one of five different mixture
  weightings.
docstr: Initializes the instance.
  
  Args:
    prior_fn: A callable returning a `tfp.distributions.Distribution` object,
      typically a `Distribution` class or factory function. This is 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 distributions 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 callable will receive keyword arguments as determined
      by `parameter_fns`.
    index_ranges: Iterable of integers. `indexes` must have the same shape as
      the bottleneck tensor, with an additional dimension at position
      `channel_axis`. The values of the `k`th channel must be in the range
      `[0, index_ranges[k])`.
    parameter_fns: Dict of strings to callables. Functions mapping `indexes`
      to each distribution parameter. For each item, `indexes` is passed to
      the callable, and the string key and return value make up one keyword
      argument to `prior_fn`.
    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.
    channel_axis: Integer or `None`. Determines the position of the channel
      axis in `indexes`. Defaults to the last dimension. If set to `None`,
      the index tensor is expected to have the same shape as the bottleneck
      tensor (only allowed when `index_ranges` has length 1).
    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_dtype: `tf.dtypes.DType`. Data type of prior and probability
      computations. Defaults to `tf.float32`.
    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_fn",":","A","callable","returning","a","`","tfp.distributions.Distribution","`","object",",","typically","a","`","Distribution","`","class","or","factory","function",".","This","is","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","distributions","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","callable","will","receive","keyword","arguments","as","determined","by","`","parameter_fns","`",".","index_ranges",":","Iterable","of","integers",".","`","indexes","`","must","have","the","same","shape","as","the","bottleneck","tensor",",","with","an","additional","dimension","at","position","`","channel_axis","`",".","The","values","of","the","`","k","`","th","channel","must","be","in","the","range","`","[","0",",","index_ranges","[","k","]",")","`",".","parameter_fns",":","Dict","of","strings","to","callables",".","Functions","mapping","`","indexes","`","to","each","distribution","parameter",".","For","each","item",",","`","indexes","`","is","passed","to","the","callable",",","and","the","string","key","and","return","value","make","up","one","keyword","argument","to","`","prior_fn","`",".","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",".","channel_axis",":","Integer","or","`","None","`",".","Determines","the","position","of","the","channel","axis","in","`","indexes","`",".","Defaults","to","the","last","dimension",".","If","set","to","`","None","`",",","the","index","tensor","is","expected","to","have","the","same","shape","as","the","bottleneck","tensor","(","only","allowed","when","`","index_ranges","`","has","length","1",")",".","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_dtype",":","`","tf.dtypes.DType","`",".","Data","type","of","prior","and","probability","computations",".","Defaults","to","`","tf.float32","`",".","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_fn,
      index_ranges,
      parameter_fns,
      coding_rank,
      channel_axis=-1,
      compression=False,
      stateless=False,
      expected_grads=False,
      tail_mass=2**-8,
      range_coder_precision=12,
      bottleneck_dtype=None,
      prior_dtype=tf.float32,
      decode_sanity_check=True,
      laplace_tail_mass=0,
  ):
      """Initializes the instance.
  
      Args:
        prior_fn: A callable returning a `tfp.distributions.Distribution` object,
          typically a `Distribution` class or factory function. This is 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 distributions 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 callable will receive keyword arguments as determined
          by `parameter_fns`.
        index_ranges: Iterable of integers. `indexes` must have the same shape as
          the bottleneck tensor, with an additional dimension at position
          `channel_axis`. The values of the `k`th channel must be in the range
          `[0, index_ranges[k])`.
        parameter_fns: Dict of strings to callables. Functions mapping `indexes`
          to each distribution parameter. For each item, `indexes` is passed to
          the callable, and the string key and return value make up one keyword
          argument to `prior_fn`.
        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.
        channel_axis: Integer or `None`. Determines the position of the channel
          axis in `indexes`. Defaults to the last dimension. If set to `None`,
          the index tensor is expected to have the same shape as the bottleneck
          tensor (only allowed when `index_ranges` has length 1).
        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_dtype: `tf.dtypes.DType`. Data type of prior and probability
          computations. Defaults to `tf.float32`.
        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 not callable(prior_fn):
          raise TypeError(
              "`prior_fn` must be a class or factory function."
          )
      for name, fn in parameter_fns.items():
          if not isinstance(name, str):
              raise TypeError("`parameter_fns` must have string keys.")
          if not callable(fn):
              raise TypeError(
                  f"`parameter_fns['{name}']` must be callable."
              )
  
      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._index_ranges = tuple(int(r) for r in index_ranges)
      if not self.index_ranges:
          raise ValueError(
              "`index_ranges` must have at least one element."
          )
      self._channel_axis = (
          None if channel_axis is None else int(channel_axis)
      )
      if self.channel_axis is None and len(self.index_ranges) > 1:
          raise ValueError(
              "`channel_axis` can't be `None` for `len(index_ranges) > 1`."
          )
      self._prior_fn = prior_fn
      self._parameter_fns = dict(parameter_fns)
      self._prior_dtype = tf.as_dtype(prior_dtype)
      self.decode_sanity_check = decode_sanity_check
  
      with self.name_scope:
          if self.compression:
              if self.channel_axis is None:
                  (index_range,) = index_ranges
                  indexes = tf.range(index_range, dtype=tf.int32)
              else:
                  indexes = [
                      tf.range(r, dtype=tf.int32) for r in index_ranges
                  ]
                  indexes = tf.meshgrid(*indexes, indexing="ij")
                  indexes = tf.stack(indexes, axis=self.channel_axis)
              self._prior = self._make_prior(indexes)
              cdf, cdf_offset = self._build_tables(
                  self.prior, range_coder_precision
              )
              self._init_compression(cdf, cdf_offset, None)
  
imports: 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
  from tensorflow_compression.python.ops import round_ops
label_desc: Use image node_id 1 to create a new ContinuousIndexedEntropyModel object from inherited base classes: continuous_base with example: obj = ContinuousIndexedEntropyModel(prior_fn, index_ranges, parameter_fns, coding_rank, channel_axis, compression, stateless, expected_grads, tail_mass, range_coder_precision, bottleneck_dtype, prior_dtype, decode_sanity_check, laplace_tail_mass)