Added extra requirements for tensorflow, tensorflow-gpu, and tensorfl…

…ow-cpu (tensorflow#2114) * Added extra requirements for tensorflow, tensorflow-gpu, and tensorflow-cpu
jrruijli · Sep 2, 2020 · 5ab03f6 · 5ab03f6
1 parent 33e2869
commit 5ab03f6
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -52,6 +52,15 @@ run the following:
 ```
 pip install tensorflow-addons
 ```
+
+To ensure you have a version of TensorFlow that is compatible with TensorFlow Addons, 
+you can specify the `tensorflow` extra requirement during install:
+
+```
+pip install tensorflow-addons[tensorflow]
+```
+
+Similar extras exist for the `tensorflow-gpu` and `tensorflow-cpu` packages.
 
 
 To use TensorFlow Addons:

diff --git a/setup.py b/setup.py
@@ -76,6 +76,8 @@ def has_ext_modules(self):
 
 
 project_name, version = get_project_name_version()
+min_tf_version = version["MIN_TF_VERSION"]
+max_tf_version = version["MAX_TF_VERSION"]
 setup(
     name=project_name,
     version=version["__version__"],
@@ -86,6 +88,15 @@ def has_ext_modules(self):
     packages=find_packages(),
     ext_modules=get_ext_modules(),
     install_requires=Path("requirements.txt").read_text().splitlines(),
+    extras_require={
+        "tensorflow": ["tensorflow>={},<{}".format(min_tf_version, max_tf_version)],
+        "tensorflow-gpu": [
+            "tensorflow-gpu>={},<{}".format(min_tf_version, max_tf_version)
+        ],
+        "tensorflow-cpu": [
+            "tensorflow-cpu>={},<{}".format(min_tf_version, max_tf_version)
+        ],
+    },
     include_package_data=True,
     zip_safe=False,
     distclass=BinaryDistribution,

diff --git a/tensorflow_addons/utils/ensure_tf_install.py b/tensorflow_addons/utils/ensure_tf_install.py
@@ -23,8 +23,7 @@
 
 import tensorflow as tf
 
-MIN_TF_VERSION = "2.2.0"
-MAX_TF_VERSION = "2.4.0"
+from tensorflow_addons.version import MIN_TF_VERSION, MAX_TF_VERSION
 
 
 def _check_tf_version():

diff --git a/tensorflow_addons/version.py b/tensorflow_addons/version.py
@@ -14,6 +14,10 @@
 # ============================================================================
 """Define TensorFlow Addons version information."""
 
+# Required TensorFlow version [min, max)
+MIN_TF_VERSION = "2.2.0"
+MAX_TF_VERSION = "2.4.0"
+
 # We follow Semantic Versioning (https://semver.org/)
 _MAJOR_VERSION = "0"
 _MINOR_VERSION = "12"