Experimental Apple Silicon (M1) support (#1924)

GPflow · Jul 5, 2022 · 4373ecc · 4373ecc
1 parent 5110164
commit 4373ecc
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 5 deletions.
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
@@ -83,5 +83,6 @@ Because GitHub's [graph of contributors](http://github.com/GPflow/GPflow/graphs/
 [@jesnie](https://github.com/jesnie)
 [@tmct](https://github.com/tmct)
 [@ltiao](https://github.com/ltiao)
+[@corwinpro](https://github.com/corwinpro)
 
 Add yourself when you first contribute to GPflow's code, tests, or documentation!
diff --git a/RELEASE.md b/RELEASE.md
@@ -61,12 +61,13 @@ This release contains contributions from:
 * <SIMILAR TO ABOVE SECTION, BUT FOR OTHER IMPORTANT CHANGES / BUG FIXES>
 * <IF A CHANGE CLOSES A GITHUB ISSUE, IT SHOULD BE DOCUMENTED HERE>
 * <NOTES SHOULD BE GROUPED PER AREA>
+* Add support for Apple Silicon Macs (`arm64`) via the `tensorflow-macos` dependency. (#1850)
 
 ## Thanks to our Contributors
 
 This release contains contributions from:
 
-jesnie
+jesnie, corwinpro
 
 
 # Release 2.5.2

diff --git a/gpflow/utilities/traversal.py b/gpflow/utilities/traversal.py
@@ -17,6 +17,7 @@
 from functools import lru_cache
 from typing import Any, Callable, Dict, Mapping, Optional, Pattern, Tuple, Type, TypeVar, Union
 
+import numpy as np
 import tensorflow as tf
 import tensorflow_probability as tfp
 from packaging.version import Version
@@ -330,7 +331,7 @@ def _first_three_elements_regexp() -> Pattern[str]:
 
 
 def _str_tensor_value(value: AnyNDArray) -> str:
-    value_str = str(value)
+    value_str = str(np.around(value, 5))
     if value.size <= 3:
         return value_str
 

diff --git a/setup.py b/setup.py
@@ -17,7 +17,9 @@
     "setuptools>=41.0.0",  # to satisfy dependency constraints
     "tabulate",
     "tensorflow-probability>=0.12.0",
-    "tensorflow>=2.4.0",
+    "tensorflow>=2.4.0; platform_system!='Darwin' or platform_machine!='arm64'",
+    # NOTE: Support of Apple Silicon MacOS platforms is in an experimental mode
+    "tensorflow-macos>=2.4.0; platform_system=='Darwin' and platform_machine=='arm64'",
     # NOTE: once we require tensorflow-probability>=0.12, we can remove our custom deepcopy handling
     "typing_extensions",
 ]

diff --git a/tests/gpflow/utilities/test_traversal.py b/tests/gpflow/utilities/test_traversal.py
@@ -328,7 +328,7 @@ def test_leaf_components_registers_variable_properties(
     for path, variable in leaf_components(module).items():
         param_name = path.split(".")[-2] + "." + path.split(".")[-1]
         assert isinstance(variable, gpflow.Parameter)
-        np.testing.assert_equal(variable.numpy(), expected_param_dicts[param_name]["value"])
+        np.testing.assert_almost_equal(variable.numpy(), expected_param_dicts[param_name]["value"])
         assert variable.trainable == expected_param_dicts[param_name]["trainable"]
         assert variable.shape == expected_param_dicts[param_name]["shape"]
 
@@ -349,7 +349,7 @@ def test_leaf_components_registers_compose_kernel_variable_properties(
         path_as_list = path.split(".")
         param_name = path_as_list[-3] + "." + path_as_list[-2] + "." + path_as_list[-1]
         assert isinstance(variable, gpflow.Parameter)
-        np.testing.assert_equal(variable.numpy(), expected_param_dicts[param_name]["value"])
+        np.testing.assert_almost_equal(variable.numpy(), expected_param_dicts[param_name]["value"])
         assert variable.trainable == expected_param_dicts[param_name]["trainable"]
         assert variable.shape == expected_param_dicts[param_name]["shape"]