Pin tensorflow_probability to work with TF 2.4.1 (#1057)

* Pin tensorflow_probability to work with TF 2.4.1

The latest version of `tensorflow_probability` is not compatible with
our version of TF:

```
ImportError: This version of TensorFlow Probability requires TensorFlow
version >= 2.5; Detected an installation of version 2.4.1. Please
upgrade TensorFlow to proceed```

Pinning & adding a test to prevent regression.

http://b/194837139

* import tensorflow

Author: rosbo

Dockerfile

@@ -62,6 +62,7 @@ RUN pip install seaborn python-dateutil dask python-igraph && \
 RUN pip install tensorflow==${TENSORFLOW_VERSION} && \
     pip install tensorflow-gcs-config==2.4.0 && \
     pip install tensorflow-addons==0.12.1 && \
+    pip install tensorflow_probability==0.12.2 && \
     /tmp/clean-layer.sh
 
 RUN apt-get install -y libfreetype6-dev && \

tests/test_tensorflow_probability.py

@@ -0,0 +1,12 @@
+import unittest
+
+# b/194837139 importing tensorflow before tfp was trigerring an error. Adding this import to prevent regression.
+import tensorflow
+import tensorflow_probability as tfp
+
+
+class TestTensorFlowProbability(unittest.TestCase):
+    def test_distribution(self):
+        tfd = tfp.distributions
+        dist = tfd.Bernoulli(logits=[-1, 1, 1])
+        self.assertEqual('Bernoulli', dist.name)