Updates for 2.8.0 release.

PiperOrigin-RevId: 427379522
Change-Id: I5346a758e69670960ab0bf4abe6e99466a0c49e7

Author: Johannes Ballé

docs/api_docs/python/tfc/ContinuousBatchedEntropyModel.md

@@ -20,7 +20,7 @@ description: Batched entropy model for continuous random variables.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L33-L421">
+  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L33-L426">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -34,9 +34,10 @@ Batched entropy model for continuous random variables.
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>tfc.ContinuousBatchedEntropyModel(
     prior=None, coding_rank=None, compression=False, stateless=False,
-    expected_grads=False, tail_mass=(2 ** -8), range_coder_precision=12, dtype=None,
-    prior_shape=None, cdf=None, cdf_offset=None, cdf_shapes=None,
-    offset_heuristic=True, quantization_offset=None, laplace_tail_mass=0
+    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,
+    laplace_tail_mass=0
 )
 </code></pre>
 
@@ -193,11 +194,11 @@ Integer. Precision passed to the range coding op.
 </td>
 </tr><tr>
 <td>
-`dtype`
+`bottleneck_dtype`
 </td>
 <td>
-`tf.dtypes.DType`. Data type of this entropy model (i.e. dtype of
-prior, decompressed values). Must be provided if `prior` is omitted.
+`tf.dtypes.DType`. Data type of bottleneck tensor.
+Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
 </td>
 </tr><tr>
 <td>
@@ -270,6 +271,13 @@ Laplace mixture for training stability. (experimental)
 
 <tr>
 <td>
+`bottleneck_dtype`
+</td>
+<td>
+Data type of the bottleneck tensor.
+</td>
+</tr><tr>
+<td>
 `cdf`
 </td>
 <td>
@@ -298,13 +306,6 @@ Whether this entropy model is prepared for compression.
 </td>
 </tr><tr>
 <td>
-`dtype`
-</td>
-<td>
-Data type of this entropy model.
-</td>
-</tr><tr>
-<td>
 `expected_grads`
 </td>
 <td>
@@ -456,7 +457,7 @@ of calling this method if you don't expect the return value to change.
 
 <h3 id="compress"><code>compress</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L336-L371">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L340-L376">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>compress(
@@ -512,7 +513,7 @@ coding unit.
 
 <h3 id="decompress"><code>decompress</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L373-L407">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L378-L412">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>decompress(
@@ -568,7 +569,7 @@ self.prior_shape`.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L409-L421">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L414-L426">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
@@ -593,7 +594,7 @@ A JSON-serializable Python dict.
 
 <h3 id="get_weights"><code>get_weights</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L339-L340">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L340-L341">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_weights()
@@ -604,7 +605,7 @@ A JSON-serializable Python dict.
 
 <h3 id="quantize"><code>quantize</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L314-L334">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L317-L338">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>quantize(
@@ -657,7 +658,7 @@ A `tf.Tensor` containing the quantized values.
 
 <h3 id="set_weights"><code>set_weights</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L342-L347">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L343-L348">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>set_weights(
@@ -733,7 +734,7 @@ The original method wrapped such that it enters the module's name scope.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L282-L312">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_batched.py#L284-L315">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__call__(

docs/api_docs/python/tfc/ContinuousIndexedEntropyModel.md

@@ -21,7 +21,7 @@ description: Indexed entropy model for continuous random variables.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L30-L408">
+  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L30-L420">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -36,7 +36,8 @@ Indexed entropy model for continuous random variables.
 <code>tfc.ContinuousIndexedEntropyModel(
     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, dtype=tf.float32, laplace_tail_mass=0
+    range_coder_precision=12, bottleneck_dtype=None, prior_dtype=tf.float32,
+    laplace_tail_mass=0
 )
 </code></pre>
 
@@ -247,11 +248,19 @@ Integer. Precision passed to the range coding op.
 </td>
 </tr><tr>
 <td>
-`dtype`
+`bottleneck_dtype`
 </td>
 <td>
-`tf.dtypes.DType`. Data type of this entropy model (i.e. dtype of
-prior, decompressed values).
+`tf.dtypes.DType`. Data type of bottleneck tensor.
+Defaults to `tf.keras.mixed_precision.global_policy().compute_dtype`.
+</td>
+</tr><tr>
+<td>
+`prior_dtype`
+</td>
+<td>
+`tf.dtypes.DType`. Data type of prior and probability
+computations. Defaults to `tf.float32`.
 </td>
 </tr><tr>
 <td>
@@ -275,6 +284,13 @@ laplace mixture for training stability. (experimental)
 
 <tr>
 <td>
+`bottleneck_dtype`
+</td>
+<td>
+Data type of the bottleneck tensor.
+</td>
+</tr><tr>
+<td>
 `cdf`
 </td>
 <td>
@@ -310,13 +326,6 @@ Whether this entropy model is prepared for compression.
 </td>
 </tr><tr>
 <td>
-`dtype`
-</td>
-<td>
-Data type of this entropy model.
-</td>
-</tr><tr>
-<td>
 `expected_grads`
 </td>
 <td>
@@ -380,6 +389,13 @@ Prior distribution, used for deriving range coding tables.
 </td>
 </tr><tr>
 <td>
+`prior_dtype`
+</td>
+<td>
+Data type of `prior`.
+</td>
+</tr><tr>
+<td>
 `prior_fn`
 </td>
 <td>
@@ -461,7 +477,7 @@ of calling this method if you don't expect the return value to change.
 
 <h3 id="compress"><code>compress</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L338-L369">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L349-L381">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>compress(
@@ -523,7 +539,7 @@ coding unit.
 
 <h3 id="decompress"><code>decompress</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L371-L397">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L383-L409">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>decompress(
@@ -577,7 +593,7 @@ dimension).
 
 <h3 id="from_config"><code>from_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L404-L408">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L416-L420">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>@classmethod</code>
@@ -591,7 +607,7 @@ Instantiates an entropy model from a configuration dictionary.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L399-L402">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L411-L414">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config()
@@ -602,7 +618,7 @@ Returns the configuration of the entropy model.
 
 <h3 id="get_weights"><code>get_weights</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L339-L340">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L340-L341">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_weights()
@@ -613,7 +629,7 @@ Returns the configuration of the entropy model.
 
 <h3 id="quantize"><code>quantize</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L320-L336">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L330-L347">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>quantize(
@@ -662,7 +678,7 @@ A `tf.Tensor` containing the quantized values.
 
 <h3 id="set_weights"><code>set_weights</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L342-L347">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_base.py#L343-L348">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>set_weights(
@@ -738,7 +754,7 @@ The original method wrapped such that it enters the module's name scope.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L277-L318">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/entropy_models/continuous_indexed.py#L286-L328">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__call__(

docs/api_docs/python/tfc/GDN.md

@@ -25,7 +25,7 @@ description: Generalized divisive normalization layer.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L31-L461">
+  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L41-L466">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -779,7 +779,7 @@ using a `keras.Metric.Mean`.
 
 <h3 id="build"><code>build</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L330-L363">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L336-L369">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>build(
@@ -868,7 +868,7 @@ one per output tensor of the layer).
 
 <h3 id="compute_output_shape"><code>compute_output_shape</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L418-L419">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L423-L424">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>compute_output_shape(
@@ -1007,7 +1007,7 @@ A layer instance.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L421-L461">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/gdn.py#L426-L466">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config() -> Dict[str, Any]

docs/api_docs/python/tfc/GDNParameter.md

@@ -17,7 +17,7 @@ description: Nonnegative parameterization as needed for GDN parameters.
 
 <table class="tfo-notebook-buttons tfo-api nocontent" align="left">
 <td>
-  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L169-L248">
+  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L186-L269">
     <img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />
     View source on GitHub
   </a>
@@ -220,7 +220,7 @@ of calling this method if you don't expect the return value to change.
 
 <h3 id="get_config"><code>get_config</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L239-L248">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L260-L269">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config() -> Dict[str, Any]
@@ -318,10 +318,12 @@ The original method wrapped such that it enters the module's name scope.
 
 <h3 id="__call__"><code>__call__</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L225-L229">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L240-L250">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
-<code>__call__() -> tf.Tensor
+<code>__call__(
+    compute_dtype=None
+) -> tf.Tensor
 </code></pre>
 
 Computes and returns the non-negative value as a `tf.Tensor`.