better m1 compatibility

Author: refraction-ray

check_all.sh

@@ -9,5 +9,5 @@ pylint tensorcircuit tests examples/*.py
 echo "pytest check"
 pytest -n 4 --cov=tensorcircuit -vv -W ignore::DeprecationWarning
 echo "sphinx check"
-cd docs && sphinx-build source build/html && sphinx-build source -D language="zh"  -D master_doc=index_cn build/html_cn
+cd docs && sphinx-build source build/html && sphinx-build source -D language="zh" build/html_cn
 echo "all checks passed, congratulation! 💐"

setup.py

@@ -1,20 +1,17 @@
 import setuptools
-import platform
 
 from tensorcircuit import __version__, __author__
-
+from tensorcircuit.utils import is_m1mac
 
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
 install_requires = ["numpy", "scipy", "tensornetwork", "networkx"]
 
-if platform.processor() != "arm":
+if not is_m1mac():
     install_requires.append("tensorflow")
     # avoid the embarassing macos M1 chip case, where the package is called tensorflow-macos
 
-# TODO(@refraction-ray): add better check_m1 function
-
 setuptools.setup(
     name="tensorcircuit",
     version=__version__,

tensorcircuit/quantum.py

@@ -37,6 +37,7 @@
 
 from .cons import backend, contractor, dtypestr, npdtype
 from .backends import get_backend  # type: ignore
+from .utils import is_m1mac
 
 Tensor = Any
 Graph = Any
@@ -1103,7 +1104,10 @@ def quimb2qop(qb_mpo: Any) -> QuOperator:
 
 
 try:
-    compiled_jit = partial(get_backend("tensorflow").jit, jit_compile=True)
+    if is_m1mac():
+        compiled_jit = lambda x: x
+    else:
+        compiled_jit = partial(get_backend("tensorflow").jit, jit_compile=True)
 
     def heisenberg_hamiltonian(
         g: Graph,

tensorcircuit/utils.py

@@ -4,6 +4,7 @@
 
 from typing import Any, Callable, Union, Sequence
 from functools import wraps
+import platform
 
 
 def return_partial(
@@ -81,3 +82,17 @@ def wrapper(*args: Any, **kws: Any) -> Any:
         return rs
 
     return wrapper
+
+
+def is_m1mac() -> bool:
+    """
+    check whether the running platform is MAC with M1 chip
+
+    :return: True for MAC M1 platform
+    :rtype: bool
+    """
+    if platform.processor() != "arm":
+        return False
+    if not platform.platform().startswith("macOS"):
+        return False
+    return True