Prepares for 2.1 release.

- Removes workarounds for TF 2.4, since they are not enough to make TFC fully
  functional. Instead, we depend on TF nightly builds.
- Removes dynamic check for TF version, and moves TF into `requirements.txt`.
  This is possible since in TF 2, the `-gpu` pip packages have been
  discontinued.
- Simplifies/adjusts installation instructions accordingly.

PiperOrigin-RevId: 362235658
Change-Id: I939d07387cdbe244493b37e659a9bbfe21ee06da

Author: Johannes Ballé

README.md

@@ -36,29 +36,31 @@ for a complete description of the classes and functions this package implements.
 ## Installation
 
 ***Note: Precompiled packages are currently only provided for Linux and
-Darwin/Mac OS. To use these packages on Windows, consider using a [TensorFlow
-Docker image](https://www.tensorflow.org/install/docker) and installing
-TensorFlow Compression using pip inside the Docker container.***
+Darwin/Mac OS and Python 3.6-3.8. To use these packages on Windows, consider
+using a [TensorFlow Docker image](https://www.tensorflow.org/install/docker) and
+installing TensorFlow Compression using pip inside the Docker container.***
 
 Set up an environment in which you can install precompiled binary Python
 packages using the `pip` command. Refer to the
 [TensorFlow installation instructions](https://www.tensorflow.org/install/pip)
 for more information on how to set up such a Python environment.
 
-The current version of TFC (v2.0) requires TensorFlow v2.4. For versions
-compatible with TensorFlow v1, see our [previous
+The current version of TensorFlow Compression requires TensorFlow 2. For
+versions compatible with TensorFlow 1, see our [previous
 releases](https://github.com/tensorflow/compression/releases).
+***Note: Because TFC currently relies on features and fixes designated for
+TF 2.5, the pip package currently depends on tf-nightly packages. Once TF 2.5 is
+released (likely in April 2021), we will resume depending on the stable version
+of TF.***
 
 ### pip
 
-To install TF and TFC via `pip`, run the following command:
+To install TFC via `pip`, run the following command:
 
 ```bash
-pip install tensorflow-gpu==2.4.* tensorflow-compression==2.0
+pip install tensorflow-compression
 ```
 
-If you don't need GPU support, you can drop the `-gpu` part.
-
 To test that the installation works correctly, you can run the unit tests with:
 
 ```bash
@@ -77,8 +79,8 @@ and then run the `pip install` command inside the Docker container, not on the
 host. For instance, you can use a command line like this:
 
 ```bash
-docker run tensorflow/tensorflow:2.4.1 bash -c \
-    "pip install tensorflow-compression==2.0 &&
+docker run tensorflow/tensorflow:nightly bash -c \
+    "pip install tensorflow-compression &&
      python -m tensorflow_compression.all_tests"
 ```
 
@@ -91,13 +93,12 @@ It seems that [Anaconda](https://www.anaconda.com/distribution/) ships its own
 binary version of TensorFlow which is incompatible with our pip package. To
 solve this, always install TensorFlow via `pip` rather than `conda`. For
 example, this creates an Anaconda environment with Python 3.6 and CUDA
-libraries, and then installs TensorFlow and TensorFlow Compression with GPU
-support:
+libraries, and then installs TensorFlow and TensorFlow Compression:
 
 ```bash
-conda create --name ENV_NAME python=3.6 cudatoolkit=10.0 cudnn
+conda create --name ENV_NAME python=3.8 cudatoolkit=10.0 cudnn
 conda activate ENV_NAME
-pip install tensorflow-gpu==2.4.* tensorflow-compression==2.0
+pip install tensorflow-compression
 ```
 
 ## Usage

docs/api_docs/python/tfc/DeepFactorized.md

@@ -190,6 +190,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `init_scale`
 </td>
 <td>

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#L149-L229">
+  <a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L147-L226">
     <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#L220-L229">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L217-L226">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_config() -> Dict[str, Any]
@@ -231,7 +231,7 @@ Returns the configuration of the `Parameter`.
 
 <h3 id="get_weights"><code>get_weights</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L49-L50">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L48-L49">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>get_weights()
@@ -242,7 +242,7 @@ Returns the configuration of the `Parameter`.
 
 <h3 id="set_weights"><code>set_weights</code></h3>
 
-<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L52-L57">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L51-L56">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>set_weights(
@@ -318,7 +318,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/layers/parameters.py#L207-L210">View source</a>
+<a target="_blank" href="https://github.com/tensorflow/compression/tree/master/tensorflow_compression/python/layers/parameters.py#L203-L207">View source</a>
 
 <pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
 <code>__call__() -> tf.Tensor

docs/api_docs/python/tfc/MonotonicAdapter.md

@@ -153,6 +153,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `name`
 </td>
 <td>

docs/api_docs/python/tfc/NoisyDeepFactorized.md

@@ -148,6 +148,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `name`
 </td>
 <td>

docs/api_docs/python/tfc/NoisyLogistic.md

@@ -148,6 +148,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `name`
 </td>
 <td>

docs/api_docs/python/tfc/NoisyLogisticMixture.md

@@ -246,6 +246,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `mixture_distribution`
 </td>
 <td>

docs/api_docs/python/tfc/NoisyMixtureSameFamily.md

@@ -245,6 +245,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `mixture_distribution`
 </td>
 <td>

docs/api_docs/python/tfc/NoisyNormal.md

@@ -148,6 +148,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `name`
 </td>
 <td>

docs/api_docs/python/tfc/NoisyNormalMixture.md

@@ -246,6 +246,13 @@ May be partially defined or unknown.
 </td>
 </tr><tr>
 <td>
+`experimental_is_sharded`
+</td>
+<td>
+`True` for distributions which parallel-sum `log_prob` across devices.
+</td>
+</tr><tr>
+<td>
 `mixture_distribution`
 </td>
 <td>