Repository: google/jax

Issue:
Build for Ubuntu (ARM-based machine)
I was trying to build jax on my Jetson Nano device. Although failed at first, I finally managed to build it by making the following change at [line 326](https://github.com/google/jax/blob/master/build/build.py#L326) in `build/build.py`:

```
[":install_xla_in_source_tree", os.getcwd()] + ["--cpu=arm"])
```
which follows the suggestion from [here](https://blog.bazel.build/2019/02/11/configurable-builds-part-1.html).
I wonder if it is useful to make this in a more general setting (maybe add a flag as an option in build) for other people to build this wonderful package on ARM-based machines. If not, I am happy to share this piece of information in case anyone having the same problem.


Discussion:
Thanks! It's great to know that JAX works on ARM! I'll add a flag to the build file to make this easier (or add a way to pass your own flags, like `--cpu` to the bazel command line.)
Thank you very much! BTW when I was importing jax it shows the following:

```
>>> from jax import random
>>> key = random.PRNGKey(0)
2019-05-28 16:38:35.228715: W external/org_tensorflow/tensorflow/core/platform/profile_utils/cpu_utils.cc:98] Failed to find bogomips in /proc/cpuinfo; cannot determine CPU frequency
2019-05-28 16:38:38.662329: F external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:175] Check failed: err == cudaSuccess || err == cudaErrorInvalidValue Unexpected CUDA error: unknown error
Aborted (core dumped)
```

I don't think it is a problem of jax itself but I wonder if there is a way to get around it. (This also happens when calling TF, but TF somehow ignores it and continue to compute on GPU.)

Thank you for your help!
Since we don't have such a device we can't easily debug it. Contributions are welcome.

The first message is a warning (`W`) and seems benign. The second message is a fatal error (`F`), but I don't know what's causing it since the error message from the CUDA libraries is not helpful at all...


Thank you so much!
@mark-fangzhou-xie 
Hi, When I follow your instructions to build jax on Nano, it shows the error:
```
Traceback (most recent call last):
  File "build/build.py", line 380, in <module>
    main()
  File "build/build.py", line 333, in main
    check_bazel_version(bazel_path, min_version="2.0.0", max_version=None)
  File "build/build.py", line 157, in check_bazel_version
    version_output = shell([bazel_path, "--bazelrc=/dev/null", "version"])
  File "build/build.py", line 47, in shell
    output = subprocess.check_output(cmd)
  File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 389, in check_output
    **kwargs).stdout
  File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 466, in run
    with Popen(*popenargs, **kwargs) as process:
  File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 769, in __init__
    restore_signals, start_new_session)
  File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 1516, in _execute_child
    raise child_exception_type(errno_num, err_msg, err_filename)
OSError: [Errno 8] Exec format error: './bazel-2.0.0-linux-x86_64'
```
How can I fix it ? Thanks
I use the command: python build/build.py --enable_cuda
> @mark-fangzhou-xie
> Hi, When I follow your instructions to build jax on Nano, it shows the error:
> 
> ```
> Traceback (most recent call last):
>   File "build/build.py", line 380, in <module>
>     main()
>   File "build/build.py", line 333, in main
>     check_bazel_version(bazel_path, min_version="2.0.0", max_version=None)
>   File "build/build.py", line 157, in check_bazel_version
>     version_output = shell([bazel_path, "--bazelrc=/dev/null", "version"])
>   File "build/build.py", line 47, in shell
>     output = subprocess.check_output(cmd)
>   File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 389, in check_output
>     **kwargs).stdout
>   File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 466, in run
>     with Popen(*popenargs, **kwargs) as process:
>   File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 769, in __init__
>     restore_signals, start_new_session)
>   File "/home/lly2014/archiconda3/lib/python3.7/subprocess.py", line 1516, in _execute_child
>     raise child_exception_type(errno_num, err_msg, err_filename)
> OSError: [Errno 8] Exec format error: './bazel-2.0.0-linux-x86_64'
> ```
> 
> How can I fix it ? Thanks
> I use the command: python build/build.py --enable_cuda

As you can see, I didn't really solve the problem myself (lack of experience in C++ and CUDA), so I really can't help you with this. 

I later built a fully-fledged Ubuntu machine and it works smoothly :)
> I was trying to build jax on my Jetson Nano device. Although failed at first, I finally managed to build it by making the following change at line 326 in build/build.py:

didn't you mange to build it? 
> > I was trying to build jax on my Jetson Nano device. Although failed at first, I finally managed to build it by making the following change at line 326 in build/build.py:
> 
> didn't you mange to build it?

The build seemed to be successful. That's true. No error came out during the building process. And I thought it was a success. But that doesn't mean it will work. If you notice, I didn't even manage to set the random number key later on after importing it. I wouldn't call that "success" to be honest. 

As for now, I completely forgot what I did at that time and don't have nano at hand. There isn't really much I could do to help you here.
Did you have any luck further? I'm trying to build it on my Jetson nano currently and I'm having issues myself haha.

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -79,17 +79,22 @@ def check_python_version(python_version):
 BAZEL_BASE_URI = "https://github.com/bazelbuild/bazel/releases/download/3.7.2/"
 BazelPackage = collections.namedtuple("BazelPackage", ["file", "sha256"])
 bazel_packages = {
-    "Linux":
+    ("Linux", "x86_64"):
         BazelPackage(
             file="bazel-3.7.2-linux-x86_64",
             sha256=
             "70dc0bee198a4c3d332925a32d464d9036a831977501f66d4996854ad4e4fc0d"),
-    "Darwin":
+    ("Linux", "aarch64"):
+        BazelPackage(
+            file="bazel-3.7.2-linux-arm64",
+            sha256=
+            "6ebd9eccbcb8f63c92a324c0c86cec11963aa9dcb914dd4718f592fdfeda9823"),
+    ("Darwin", "x86_64"):
         BazelPackage(
             file="bazel-3.7.2-darwin-x86_64",
             sha256=
             "80c82e93a12ba30021692b11c78007807e82383a673be1602573b944beb359ab"),
-    "Windows":
+    ("Windows", "x86_64"):
         BazelPackage(
             file="bazel-3.7.2-windows-x86_64.exe",
             sha256=
@@ -99,7 +104,7 @@ def check_python_version(python_version):
 
 def download_and_verify_bazel():
   """Downloads a bazel binary from Github, verifying its SHA256 hash."""
-  package = bazel_packages.get(platform.system())
+  package = bazel_packages.get((platform.system(), platform.machine()))
   if package is None:
     return None
 
diff --git a/build/build_wheel.py b/build/build_wheel.py
--- a/build/build_wheel.py
+++ b/build/build_wheel.py
@@ -216,12 +216,16 @@ def prepare_wheel(sources_path):
 
 def build_wheel(sources_path, output_path):
   """Builds a wheel in `output_path` using the source tree in `sources_path`."""
-  platform_name = {
-    "Linux": "manylinux2010",
-    "Darwin": "macosx_10_9",
-    "Windows": "win",
-  }[platform.system()]
-  cpu_name = "amd64" if platform.system() == "Windows" else "x86_64"
+  if platform.system() == "Windows":
+    cpu_name = "amd64"
+    platform_name = "win"
+  else:
+    platform_name, cpu_name = {
+      ("Linux", "x86_64"): ("manylinux2010", "x86_64"),
+      ("Linux", "aarch64"): ("manylinux2014", "aarch64"),
+      ("Darwin", "x86_64"): ("macosx_10_9", "x86_64"),
+      ("Darwin", "arm64"): ("macosx_11_0", "arm64"),
+    }[(platform.system(), platform.machine())]
   python_tag_arg = (f"--python-tag=cp{sys.version_info.major}"
                     f"{sys.version_info.minor}")
   platform_tag_arg = f"--plat-name={platform_name}_{cpu_name}"


Test patch: