Fix tensorflow requirements for ARM Macs (#142)

Fix gh-141. Signed-off-by: Juan Luis Cano Rodríguez <juan_luis_cano@mckinsey.com>
kedro-org · Mar 26, 2023 · e4864d1 · e4864d1
1 parent a5b1c13
commit e4864d1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/kedro-datasets/setup.py b/kedro-datasets/setup.py
@@ -86,7 +86,9 @@ def _collect_requirements(requires):
     "tensorflow.TensorflowModelDataset": [
         # currently only TensorFlow V2 supported for saving and loading.
         # V1 requires HDF5 and serialises differently
-        "tensorflow~=2.0"
+        "tensorflow~=2.0; platform_system != 'Darwin' or platform_machine != 'arm64'",
+        # https://developer.apple.com/metal/tensorflow-plugin/
+        "tensorflow-macos~=2.0; platform_system == 'Darwin' and platform_machine == 'arm64'",
     ]
 }
 yaml_require = {"yaml.YAMLDataSet": [PANDAS, "PyYAML>=4.2, <7.0"]}

diff --git a/kedro-datasets/test_requirements.txt b/kedro-datasets/test_requirements.txt
@@ -55,7 +55,8 @@ snowflake-snowpark-python~=1.0.0; python_version == '3.8'
 SQLAlchemy~=1.2
 tables~=3.6.0; platform_system == "Windows" and python_version<'3.9'
 tables~=3.6; platform_system != "Windows"
-tensorflow~=2.0
+tensorflow-macos~=2.0; platform_system == "Darwin" and platform_machine == "arm64"
+tensorflow~=2.0; platform_system != "Darwin" or platform_machine != "arm64"
 triad>=0.6.7, <1.0
 trufflehog~=2.1
 xlsxwriter~=1.0