From 463f5f8d7de9b1d9516f92f0629e144d05194b49 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Yan=20Facai=20=28=E9=A2=9C=E5=8F=91=E6=89=8D=29?=
 <facai.yan@gmail.com>
Date: Wed, 4 Sep 2019 09:33:49 +0800
Subject: [PATCH] BLD: add tf 2.0 dependency for addons (#474)

* BLD: add tf dependency for addons
* DOC: remove note in readme
* BLD: support tfa-nightly
* BLD: pin to 2.0.0-rc0
---
 README.md                     |  1 -
 setup.py                      | 35 ++++++++++++++++-----
 tensorflow_addons/__init__.py | 58 +++--------------------------------
 3 files changed, 33 insertions(+), 61 deletions(-)

diff --git a/README.md b/README.md
index a801dcc254..ecabd4ae65 100644
--- a/README.md
+++ b/README.md
@@ -47,7 +47,6 @@ To install the latest version, run the following:
 pip install tensorflow-addons
 ```
  
-**Note:** You will also need [`tensorflow==2.0.0-beta1`](https://www.tensorflow.org/beta) installed.
 
 To use addons:
 
diff --git a/setup.py b/setup.py
index f5109e43a7..5d6d3cba40 100644
--- a/setup.py
+++ b/setup.py
@@ -29,6 +29,7 @@
 from __future__ import print_function
 
 import os
+import platform
 import sys
 
 from datetime import datetime
@@ -39,6 +40,17 @@
 
 DOCLINES = __doc__.split('\n')
 
+TFA_NIGHTLY = 'tfa-nightly'
+TFA_RELEASE = 'tensorflow-addons'
+
+if '--nightly' in sys.argv:
+    project_name = TFA_NIGHTLY
+    nightly_idx = sys.argv.index('--nightly')
+    sys.argv.pop(nightly_idx)
+else:
+    project_name = TFA_RELEASE
+
+# Version
 version = {}
 base_dir = os.path.dirname(os.path.abspath(__file__))
 with open(os.path.join(base_dir, "tensorflow_addons", "version.py")) as fp:
@@ -46,17 +58,26 @@
     exec(fp.read(), version)
     # yapf: enable
 
+if project_name == TFA_NIGHTLY:
+    version['__version__'] += datetime.strftime(datetime.today(), "%Y%m%d")
+
+# Dependencies
 REQUIRED_PACKAGES = [
     'six >= 1.10.0',
 ]
 
-if '--nightly' in sys.argv:
-    project_name = 'tfa-nightly'
-    nightly_idx = sys.argv.index('--nightly')
-    sys.argv.pop(nightly_idx)
-    version['__version__'] += datetime.strftime(datetime.today(), "%Y%m%d")
-else:
-    project_name = 'tensorflow-addons'
+if project_name == TFA_RELEASE:
+    # TODO: remove if-else condition when tf supports package consolidation.
+    if platform.system() == 'Linux':
+        REQUIRED_PACKAGES.append('tensorflow-gpu == 2.0.0-rc0')
+    else:
+        REQUIRED_PACKAGES.append('tensorflow == 2.0.0-rc0')
+elif project_name == TFA_NIGHTLY:
+    # TODO: remove if-else condition when tf-nightly supports package consolidation.
+    if platform.system() == 'Linux':
+        REQUIRED_PACKAGES.append('tf-nightly-gpu-2.0-preview')
+    else:
+        REQUIRED_PACKAGES.append('tf-nightly-2.0-preview')
 
 
 class BinaryDistribution(Distribution):
diff --git a/tensorflow_addons/__init__.py b/tensorflow_addons/__init__.py
index 2a9f4df814..76b48a8940 100644
--- a/tensorflow_addons/__init__.py
+++ b/tensorflow_addons/__init__.py
@@ -17,59 +17,6 @@
 from __future__ import division
 from __future__ import print_function
 
-# We need to put some imports inside a function call below, and the function
-# call needs to come before the *actual* imports that populate the
-# tensorflow_probability namespace. Hence, we disable this lint check throughout
-# the file.
-#
-
-
-# Ensure TensorFlow is importable and its version is sufficiently recent. This
-# needs to happen before anything else, since the imports below will try to
-# import tensorflow, too.
-def _ensure_tf_install():
-    """Attempt to import tensorflow, and ensure its version is sufficient.
-
-    Raises:
-      ImportError: if either tensorflow is not importable or its version is
-      inadequate.
-    """
-    try:
-        import tensorflow as tf
-    except ImportError:
-        # Print more informative error message, then reraise.
-        print("\n\nFailed to import TensorFlow. Please note that TensorFlow is"
-              " not installed by default when you install TensorFlow Addons."
-              " This is so that users can decide whether to install the"
-              " GPU-enabled TensorFlow package. To use TensorFlow Addons,"
-              " please install the most recent version of TensorFlow, by"
-              " following instructions at https://tensorflow.org/install.\n\n")
-        raise
-
-    import distutils.version
-
-    #
-    # Update this whenever we need to depend on a newer TensorFlow release.
-    #
-    required_tensorflow_version = "2"
-
-    if (distutils.version.LooseVersion(tf.__version__) <
-            distutils.version.LooseVersion(required_tensorflow_version)):
-        raise ImportError(
-            "This version of TensorFlow Addons requires TensorFlow "
-            "version >= {required}; Detected an installation of version "
-            "{present}. Please upgrade TensorFlow to proceed.".format(
-                required=required_tensorflow_version, present=tf.__version__))
-
-
-_ensure_tf_install()
-
-# Cleanup symbols to avoid polluting namespace.
-del _ensure_tf_install
-del absolute_import
-del division
-del print_function
-
 # Local project imports
 from tensorflow_addons import activations
 from tensorflow_addons import callbacks
@@ -83,3 +30,8 @@ def _ensure_tf_install():
 from tensorflow_addons import text
 
 from tensorflow_addons.version import __version__
+
+# Cleanup symbols to avoid polluting namespace.
+del absolute_import
+del division
+del print_function