search

jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.31k stars 2.78k forks source link

Provide wheels for macOS ARM #5501

Closed ericmjl closed 2 years ago

ericmjl commented 3 years ago

Hi all,

I was digging around to see what might need to happen to allow JAX to work on Apple Silicon. Knowing that JAX gets compiled to XLA, my guess here is that XLA would need to be made Apple Silicon-compatible first before JAX could run on it. May I ask, do you all know if there are plans on the XLA team to make that happen, or is it being ignored completely? (Knowing the answer can help me make some decisions on how I should set up my development environment mostly.)

Cheers, Eric

8bitmp3 commented 3 years ago

Check out this post https://github.com/google/jax/discussions/5084 by @hawkinsp (cc @rxwei)

hawkinsp commented 3 years ago

Targeting the M1's ARM CPU shouldn't be difficult.

XLA already supports AArch64 and has done for a long time. I suspect that https://github.com/tensorflow/tensorflow/pull/45404 already did most of the work to adapt XLA to build on the M1 and all that is left is a few small changes to the .bazelrc file that JAX's build.py script generates, analogous to the changes in that TF PR.

That said, I don't have access to any M1 hardware, so this is in the "contributions welcome" category.

Targeting the GPU or the Neural Engine is likely a lot more difficult. For GPU, one would probably need to target Metal (probably doable, but not trivial), and I'm unsure how we could target the Neural Engine at this time.

jotsif commented 3 years ago

Hi! I have tried to get jaxlib working now on my apple m1, and have managed to build it with some minor changes according to tensorflow/tensorflow#45404 (branch on https://github.com/jotsif/jax/tree/jax_for_darwin_arm64).

However when loading jaxlib in python3 (3.9.2) I get the import error ImportError: dlopen(/opt/homebrew/lib/python3.9/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: _LLVMInitializeAArch64AsmPrinter

Seems to be a known issue from comment here https://github.com/tensorflow/tensorflow/pull/45404#issuecomment-760484477 and here https://github.com/elixir-nx/nx/issues/217#issuecomment-781351038.

I am investigating further but if anyone has any tip that would be welcome.

jotsif commented 3 years ago

Update: this TF PR will probably fix this issue https://github.com/tensorflow/tensorflow/pull/47594

mattjj commented 3 years ago

@jotsif thanks for that information, that's really helpful!

jotsif commented 3 years ago

Confirmed that it works 🎉 . Built with Bazel master, and https://github.com/freedomtan/tensorflow/tree/bazel_native_build_on_m1

>>> import platform
>>> platform.uname()
uname_result(system='Darwin', node='Josefs-MBP-2.lan', release='20.3.0', version='Darwin Kernel Version 20.3.0: Thu Jan 21 00:06:51 PST 2021; root:xnu-7195.81.3~1/RELEASE_ARM64_T8101', machine='arm64')
>>> from jax.lib import xla_client as xc
>>> xops = xc.ops
>>> c = xc.XlaBuilder("simple_scalar")
>>> param_shape = xc.Shape.array_shape(np.dtype(np.float32), ())
>>> x = xops.Parameter(c, 0, param_shape)
>>> y = xops.Sin(x)
>>> computation = c.Build()
>>> cpu_backend = xc.get_local_backend("cpu")
2021-03-07 09:26:16.684549: W external/org_tensorflow/tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
>>> compiled_computation = cpu_backend.compile(computation)
>>> host_input = np.array(3.0, dtype=np.float32)
>>> device_input = cpu_backend.buffer_from_pyval(host_input)
>>> device_out = compiled_computation.execute([device_input ,])
>>> device_out[0].to_py()
array(0.14112, dtype=float32)

@mattjj @hawkinsp If you want a PR I would be happy to create one, but maybe it makes more sense to wait until bazel has released a working native arm64 build and tensorflow have the necessary code in master.

erwincoumans commented 3 years ago

https://github.com/freedomtan/tensorflow/tree/bazel_native_build_on_m1

Congrats! Do you mind sharing a jax/jaxlib wheel for M1?

akbir commented 3 years ago

Hi @jotsif -

Bazel 4.1 works natively with arm64 - https://github.com/bazelbuild/bazel/issues/13099 and TF has necessary code in master.

If you have an example branch, i'm happy to help towards a PR now!

hawkinsp commented 3 years ago

@akbir I think we still need to wait for Bazel to actually release 4.1. But that should be soon I think! At that point we can probably just bump the Bazel dependency to 4.1 and hopefully everything should work on Mac ARM.

We can look into releasing Mac ARM wheels as well, although we don't yet have a way to test them (we personally do not have Mac ARM hardware yet), which gives me some pause.

hawkinsp commented 3 years ago

I believe now that at head jaxlib will build from source on a Mac M1 and pretty much everything works (*). We still don't have a great way to provide pre-built wheels yet, but hopefully this is enough to unblock everyone!

(*) It's a bit annoying to install jax still, mostly because there aren't any prebuilt scipy wheels, so you'll have to build scipy yourself to build jaxlib or use jax. I followed these instructions https://github.com/scipy/scipy/issues/13409#issuecomment-824919738 which worked for me.

akbir commented 3 years ago

Hi @hawkinsp - tried following this but hit the following error when running build.py (included full logs at the bottom).

ERROR: Analysis of target '//build:build_wheel' failed; build aborted: Analysis of target '@local_config_cc//:toolchain' failed

Unsure why toolchain isn't working an error is discussed here: https://github.com/bazelbuild/bazel/issues/13099#issuecomment-846062829.

I've also run build against bazel 4.1.0.rc5 and still get the same error.

Full error logs:

Bazel binary path: ./bazel-4.1.0rc4-darwin-arm64
Python binary path: /Users/akbirkhan/jax/venv/bin/python
Python version: 3.9
MKL-DNN enabled: yes
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
./bazel-4.1.0rc4-darwin-arm64 run --verbose_failures=true --config=short_logs --config=mkl_open_source_only :build_wheel -- --output_path=/Users/akbirkhan/jax/dist
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /Users/akbirkhan/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /Users/akbirkhan/jax/.bazelrc:
  Inherited 'build' options: --repo_env PYTHON_BIN_PATH=/Users/akbirkhan/jax/venv/bin/python --action_env=PYENV_ROOT --python_path=/Users/akbirkhan/jax/venv/bin/python --repo_env TF_NEED_CUDA=0 --action_env TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 --repo_env TF_NEED_ROCM=0 --action_env TF_ROCM_AMDGPU_TARGETS=gfx803,gfx900,gfx906,gfx1010 --distinct_host_configuration=false -c opt --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true --spawn_strategy=standalone --strategy=Genrule=standalone --enable_platform_specific_config
INFO: Found applicable config definition build:short_logs in file /Users/akbirkhan/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:mkl_open_source_only in file /Users/akbirkhan/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /Users/akbirkhan/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /Users/akbirkhan/jax/.bazelrc: --copt=-Wno-sign-compare --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading: 
Loading: 0 packages loaded
WARNING: Download from http://mirror.tensorflow.org/github.com/tensorflow/runtime/archive/3f4cd5e8a34eb2179537b8f71b1484bb0d26701f.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found
DEBUG: /private/var/tmp/_bazel_akbirkhan/30d03974cb2cfe2f871935b707b3981a/external/tf_runtime/third_party/cuda/dependencies.bzl:51:10: The following command will download NVIDIA proprietary software. By using the software you agree to comply with the terms of the license agreement that accompanies the software. If you do not agree to the terms of the license agreement, do not use the software.
Analyzing: target //build:build_wheel (0 packages loaded, 0 targets configured)
ERROR: /private/var/tmp/_bazel_akbirkhan/30d03974cb2cfe2f871935b707b3981a/external/local_config_cc/BUILD:48:19: in cc_toolchain_suite rule @local_config_cc//:toolchain: cc_toolchain_suite '@local_config_cc//:toolchain' does not contain a toolchain for cpu 'darwin_arm64'
DEBUG: Rule 'io_bazel_rules_docker' indicated that a canonical reproducible form can be obtained by modifying arguments shallow_since = "1556410077 -0400"
DEBUG: Repository io_bazel_rules_docker instantiated at:
  /Users/akbirkhan/jax/WORKSPACE:34:10: in <toplevel>
  /private/var/tmp/_bazel_akbirkhan/30d03974cb2cfe2f871935b707b3981a/external/org_tensorflow/tensorflow/workspace0.bzl:108:34: in workspace
  /private/var/tmp/_bazel_akbirkhan/30d03974cb2cfe2f871935b707b3981a/external/bazel_toolchains/repositories/repositories.bzl:37:23: in repositories
Repository rule git_repository defined at:
  /private/var/tmp/_bazel_akbirkhan/30d03974cb2cfe2f871935b707b3981a/external/bazel_tools/tools/build_defs/repo/git.bzl:199:33: in <toplevel>
ERROR: Analysis of target '//build:build_wheel' failed; build aborted: Analysis of target '@local_config_cc//:toolchain' failed
INFO: Elapsed time: 0.155s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded, 0 targets configured)
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully (0 packages loaded, 0 targets configured),
subprocess.CalledProcessError: Command '['./bazel-4.1.0rc4-darwin-arm64', 'run', '--verbose_failures=true', '--config=short_logs', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/Users/akbirkhan/jax/dist']' returned non-zero exit status 1.
Noahyt commented 3 years ago

Getting same error as @akbir .

ERROR: /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/local_config_cc/BUILD:48:19: in cc_toolchain_suite rule @local_config_cc//:toolchain: cc_toolchain_suite '@local_config_cc//:toolchain' does not contain a toolchain for cpu 'darwin_arm64'
INFO: Repository com_google_absl instantiated at:
  /Users/noah/harvard/2021/network_repair/jax/WORKSPACE:30:10: in <toplevel>
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/tensorflow/workspace2.bzl:1090:28: in workspace
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/tensorflow/workspace2.bzl:56:9: in _initialize_third_party
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/third_party/absl/workspace.bzl:12:20: in repo
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/third_party/repo.bzl:112:21: in tf_http_archive
Repository rule _tf_http_archive defined at:
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/third_party/repo.bzl:65:35: in <toplevel>
Analyzing: target //build:build_wheel (35 packages loaded, 264 targets configured)
INFO: Repository cython instantiated at:
  /Users/noah/harvard/2021/network_repair/jax/WORKSPACE:30:10: in <toplevel>
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/tensorflow/workspace2.bzl:1097:21: in workspace
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/tensorflow/workspace2.bzl:845:20: in _tf_repositories
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/third_party/repo.bzl:112:21: in tf_http_archive
Repository rule _tf_http_archive defined at:
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/org_tensorflow/third_party/repo.bzl:65:35: in <toplevel>
INFO: Repository pocketfft instantiated at:
  /Users/noah/harvard/2021/network_repair/jax/WORKSPACE:24:10: in <toplevel>
  /Users/noah/harvard/2021/network_repair/jax/third_party/pocketfft/workspace.bzl:20:17: in repo
Repository rule http_archive defined at:
  /private/var/tmp/_bazel_noah/2c44e3012980c1c9cbcf3acb09acecec/external/bazel_tools/tools/build_defs/repo/http.bzl:336:31: in <toplevel>
ERROR: Analysis of target '//build:build_wheel' failed; build aborted: Analysis of target '@local_config_cc//:toolchain' failed
INFO: Elapsed time: 10.071s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (35 packages loaded, 264 targets configured)
ERROR: Build failed. Not running target
Noahyt commented 3 years ago

FYI for those looking for a quick and dirty workaround, you can install jax and jaxlib using pip a miniconda environment running in Rosetta 2. The most recent versions of jax and jaxlib don't work (giving an error like "zsh: illegal hardware instruction"), So I ended up using jax == 0.2.10 and jaxlib==0.1.60.

hawkinsp commented 3 years ago

@akbir Do you have a working XCode installation including the command line tools? https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source

I don't think the error you are seeing is related to the Bazel version. We're currently pinning Bazel 4.1.0rc4 because that was the newest version last week. If you like you can try a different Bazel version, but 4.1.0rc4 worked for me. The easiest way to do that is to install Bazel yourself and pass --bazel_path=/somewhere/bazel to the build.py command line.

@Noahyt jaxlib 0.1.62 and newer on x86 use AVX, which Rosetta does not support (https://github.com/google/jax/blob/master/CHANGELOG.md#jaxlib-0162-march-9-2021). All recent x86 CPUs support AVX and have for a long time. We don't intend to ship wheels without AVX, although if you like you can build a jaxlib from source that does not require AVX. But I don't think we should worry about that too much, since we want a native ARM version, anyway.

akbir commented 3 years ago

Got this to finally build!! Thank you @hawkinsp

For others, I installed Xcode (not just command-line interface) and Bazel.

Also updated .bazelversion to 4.1.0 (should this be updated in the repo?) 

sudo xcode-select -s /Applications/Xcode.app/Contents/Developer
sudo xcodebuild -license
bazel sync --configure

python build.py --bazel_path=/somewhere/bazel
hawkinsp commented 3 years ago

@akbir If you want to send a PR that updates the Bazel version for Mac ARM, that sounds great! The version of bazel is chosen here: https://cs.opensource.google/jax/jax/+/master:build/build.py;drc=dacf31f2020175181014745cdabc240a10031227;l=119

akbir commented 3 years ago

would love to!

quick question - why does jax also specify the version here: https://github.com/google/jax/blob/master/.bazelversion ?

hawkinsp commented 3 years ago

@akbir If I remember correctly, that's for folks using Bazelisk (https://github.com/bazelbuild/bazelisk). I don't know if it's possible to specify a separate version for Mac arm via Bazelisk.

We can probably upgrade to 4.1.0 for all platforms, but let's not do that right away. So fixing build.py is probably enough.

apaszke commented 3 years ago

Going through the comments it looks like we've done all we could on the JAX side, and the remaining difficulties are in installing SciPy. Is this ready to be closed, or are there any outstanding issues?

akbir commented 3 years ago

Outstanding issue - is having wheels for darwin-arm64, which is blocked by scipy.

Not sure how Jax organises issues. But that's the gist^

hawkinsp commented 3 years ago

It seems from the scipy issue that they are currently targeting 1.7.1 for Mac ARM wheels.

The remaining action item on the JAX side is then (eventually) to build Mac M1 wheels as part of our release process.

reshinthadithyan commented 3 years ago

Hello, @hawkinsp. Is there any rough timeline onto when the next release is? Update : I was able to build the wheel.

hawkinsp commented 3 years ago

7254 allows us to cross-compile Mac ARM wheels on an x86 Mac machine. Unfortunately we have no way to test them, not having any Mac ARM hardware of our own or any way to emulate it, so while we can certainly release the resulting wheels as a service to the community, there will be no guarantees that they will work in any given release. We'd have to rely on the community to try them out and report bugs.

I have manually verified at least the current version of jaxlib works on a (borrowed) Macbook M1.

Because the wheels will be untested by us, for now our plan is to add a warning when JAX is imported on a Mac M1 machine that says something along the lines of "JAX on Mac ARM machines is experimental and community supported, see this github issue in the event of any problems" until such time as we can perform testing either on our own machines or through Github CI.

dcxSt commented 3 years ago

Hi, I'm using a device with the M1 chip and have been following this thread very attentively although my understanding is very basic. I've been relentlessly trying to import Jax for a few days now, most recently I switched to @hawkinsp 's PR #7254 and attempted to build jaxlib with python build/build.py (using python 3.8.8, and having uninstalled bazel) however it didn't work, this is the output : 


     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\

Starting local Bazel server and connecting to it...
Bazel binary path: ./bazel-3.7.2-darwin-x86_64
Python binary path: /opt/homebrew/anaconda3/bin/python
Python version: 3.8
NumPy version: 1.19.5
SciPy version: 1.6.2
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
./bazel-3.7.2-darwin-x86_64 run --verbose_failures=true --config=short_logs --config=avx_posix --config=mkl_open_source_only :build_wheel -- --output_path=/Users/steve/Documents/code/jax/dist --cpu=x86_64
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /Users/steve/Documents/code/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /Users/steve/Documents/code/jax/.bazelrc:
  Inherited 'build' options: --repo_env PYTHON_BIN_PATH=/opt/homebrew/anaconda3/bin/python --action_env=PYENV_ROOT --python_path=/opt/homebrew/anaconda3/bin/python --repo_env TF_NEED_CUDA=0 --action_env TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 --repo_env TF_NEED_ROCM=0 --action_env TF_ROCM_AMDGPU_TARGETS=gfx803,gfx900,gfx906,gfx1010 -c opt --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true --spawn_strategy=standalone --strategy=Genrule=standalone --enable_platform_specific_config --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /Users/steve/Documents/code/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:avx_posix in file /Users/steve/Documents/code/jax/.bazelrc: --copt=-mavx --host_copt=-mavx
INFO: Found applicable config definition build:mkl_open_source_only in file /Users/steve/Documents/code/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /Users/steve/Documents/code/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /Users/steve/Documents/code/jax/.bazelrc: --copt=-Wno-sign-compare --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading:
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Analyzing: target //build:build_wheel (1 packages loaded, 0 targets configured)
Analyzing: target //build:build_wheel (12 packages loaded, 12 targets configured)
DEBUG: Rule 'io_bazel_rules_docker' indicated that a canonical reproducible form can be obtained by modifying arguments shallow_since = "1596824487 -0400"
DEBUG: Repository io_bazel_rules_docker instantiated at:
  /Users/steve/Documents/code/jax/WORKSPACE:34:10: in <toplevel>
  /private/var/tmp/_bazel_steve/c238ce85f8bfae0489a1104fd56ad209/external/org_tensorflow/tensorflow/workspace0.bzl:108:34: in workspace
  /private/var/tmp/_bazel_steve/c238ce85f8bfae0489a1104fd56ad209/external/bazel_toolchains/repositories/repositories.bzl:35:23: in repositories
Repository rule git_repository defined at:
  /private/var/tmp/_bazel_steve/c238ce85f8bfae0489a1104fd56ad209/external/bazel_tools/tools/build_defs/repo/git.bzl:199:33: in <toplevel>
Analyzing: target //build:build_wheel (14 packages loaded, 12 targets configured)
Analyzing: target //build:build_wheel (27 packages loaded, 144 targets configured)
Analyzing: target //build:build_wheel (42 packages loaded, 272 targets configured)
Analyzing: target //build:build_wheel (97 packages loaded, 1923 targets configured)
Analyzing: target //build:build_wheel (153 packages loaded, 4414 targets configured)
Analyzing: target //build:build_wheel (164 packages loaded, 7566 targets configured)
Analyzing: target //build:build_wheel (164 packages loaded, 7623 targets configured)
INFO: Analyzed target //build:build_wheel (171 packages loaded, 14292 targets configured).
INFO: Found 1 target...
[1 / 828] [Prepa] BazelWorkspaceStatusAction stable-status.txt
[73 / 1,034] Compiling com_google_protobuf/src/google/protobuf/compiler/cpp/cpp_file.cc; 2s local ... (8 actions, 7 running)
[110 / 1,034] Compiling com_google_protobuf/src/google/protobuf/compiler/objectivec/objectivec_helpers.cc; 2s local ... (8 actions, 7 running)
[181 / 1,270] Compiling com_github_grpc_grpc/src/compiler/python_generator.cc; 4s local ... (8 actions, 7 running)
[264 / 1,270] Compiling llvm-project/llvm/lib/Support/ItaniumManglingCanonicalizer.cpp; 0s local ... (8 actions, 7 running)
[333 / 1,281] Compiling llvm-project/mlir/tools/mlir-tblgen/OpFormatGen.cpp; 5s local ... (8 actions, 7 running)
[421 / 1,394] Compiling llvm-project/llvm/lib/MC/WasmObjectWriter.cpp; 2s local ... (8 actions, 7 running)
[506 / 1,445] Compiling llvm-project/llvm/lib/DebugInfo/CodeView/StringsAndChecksums.cpp; 1s local ... (8 actions, 7 running)
ERROR: /private/var/tmp/_bazel_steve/c238ce85f8bfae0489a1104fd56ad209/external/llvm-project/mlir/BUILD:155:18: TdGenerate external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h.inc failed (Illegal instruction): mlir-tblgen failed: error executing command
  (cd /private/var/tmp/_bazel_steve/c238ce85f8bfae0489a1104fd56ad209/execroot/__main__ && \
  exec env - \
  bazel-out/darwin-opt-exec-50AE0418/bin/external/llvm-project/mlir/mlir-tblgen --gen-attrdef-decls external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.td -I external/llvm-project/mlir/include -I bazel-out/darwin-opt/bin/external/llvm-project/mlir/include -I external/llvm-project/ -I bazel-out/darwin-opt/bin/external/llvm-project/ -I external/llvm-project/mlir/include/mlir/IR -I bazel-out/darwin-opt/bin/external/llvm-project/mlir/include/mlir/IR -o bazel-out/darwin-opt/bin/external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h.inc)
Execution platform: @local_execution_config_platform//:platform
Target //build:build_wheel failed to build
INFO: Elapsed time: 147.032s, Critical Path: 10.15s
INFO: 476 processes: 9 internal, 467 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
b''
Traceback (most recent call last):
  File "build/build.py", line 587, in <module>
    main()
  File "build/build.py", line 582, in main
    shell(command)
  File "build/build.py", line 52, in shell
    output = subprocess.check_output(cmd)
  File "/opt/homebrew/anaconda3/lib/python3.8/subprocess.py", line 415, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/opt/homebrew/anaconda3/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-3.7.2-darwin-x86_64', 'run', '--verbose_failures=true', '--config=short_logs', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/Users/steve/Documents/code/jax/dist', '--cpu=x86_64']' returned non-zero exit status 1.

I've been following this guid for how to build jaxlib from source. Also some potential useful information (I have no idea): running platform.machine() in python returns "x86_64" which I'm confused about because I thought M1 uses ARM.

Unfortunately the old version of JAX that does import without complaint doesn't support differentiating through discrete Fourier transforms jnp.fft.fft, that is jax==0.2.10 and jaxlib==0.1.60

Hope this is helpful! I would really appreciate any guidance / help whatsoever. Thank you @hawkinsp for working on this issue! It's impossible for me to get any work done without JAX

akbir commented 3 years ago

Hi @dcxSt!

So a couple of thing's I'd note, if you want to build Jax on your m1.

1) Your python interpreter is using rosetta (it thinks its running on intel x86 as opposed to the m1 arm64). I'm not sure how you installed it but you'll need the native m1 version. This can be easily done using brew. I'm not well versed using anaconda but i'm sure theres a channel somewhere for it.

2) You're also building using an old version of bazel - don't use the one provided in Jax - @hawkinsp mentions that here

Alternatively after #7268 - if you have the correct python interpreter i think pip install jax should do the trick (can check later if this helps).

hawkinsp commented 3 years ago

@dcxSt @akbir is correct. You are building on a Python x86-64 interpreter. You have two choices:

I recommend using a native version if you can!

dcxSt commented 3 years ago

Hallelujah it works! Thank you @akbir thank you @hawkinsp you are my saviours!

For others, here's what I did: Installed python for arm64 (first uninstall anaconda with brew uninstall anaconda)

brew install --cask miniforge
conda init zsh
conda activate
conda install numpy scipy scikit-learn

Build jaxlib from source by cloning the jax repository, and install bazel 4

git clone https://github.com/google/jax 
brew install bazel

Then enter the jax directory cd jax and build jaxlib from source, this takes a while to run (~8 mins)

python build/build.py
pip install dist/*.whl  # installs jaxlib (includes XLA)

Then install jax with either pip install -e . or pip install jax And it works!

Edit: I forgot to mention that I installed xcode before doing any of this, the whole thing not just command line tools.

moellenh commented 3 years ago

Thanks to the above instruction, I have also managed to build jaxlib on M1.

But when importing jax, I get the following strange error (importing jaxlib seems to work fine).

Python 3.9.6 | packaged by conda-forge | (default, Jul 11 2021, 03:35:11) 
[Clang 11.1.0 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
/Users/thomas/Documents/jax/jax/lib/__init__.py:31: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/thomas/Documents/jax/jax/__init__.py", line 37, in <module>
    from . import config as _config_module
  File "/Users/thomas/Documents/jax/jax/config.py", line 18, in <module>
    from jax._src.config import config
  File "/Users/thomas/Documents/jax/jax/_src/config.py", line 26, in <module>
    from jax import lib
  File "/Users/thomas/Documents/jax/jax/lib/__init__.py", line 70, in <module>
    from jaxlib import cpu_feature_guard
ImportError: dlopen(/Users/thomas/miniforge3/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so, 2): no suitable image found.  Did find:
    /Users/thomas/miniforge3/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so: mach-o, but wrong architecture
    /Users/thomas/miniforge3/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so: mach-o, but wrong architecture

Any ideas what could be the issue here?

hawkinsp commented 3 years ago

It looks like you ended up building an x86 wheel. Can you share a description of how you built it?

On Fri, Jul 23, 2021, 7:29 PM Thomas Möllenhoff @.***> wrote:

Thanks to the above instruction, I have also managed to build jaxlib on M1.

But when importing jax, I get the following strange error (importing jaxlib seems to work fine).

Python 3.9.6 | packaged by conda-forge | (default, Jul 11 2021, 03:35:11) [Clang 11.1.0 ] on darwin Type "help", "copyright", "credits" or "license" for more information.

import jax /Users/thomas/Documents/jax/jax/lib/init.py:31: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems. warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. " Traceback (most recent call last): File "", line 1, in File "/Users/thomas/Documents/jax/jax/init.py", line 37, in from . import config as _config_module File "/Users/thomas/Documents/jax/jax/config.py", line 18, in from jax._src.config import config File "/Users/thomas/Documents/jax/jax/_src/config.py", line 26, in from jax import lib File "/Users/thomas/Documents/jax/jax/lib/init.py", line 70, in from jaxlib import cpu_feature_guard ImportError: dlopen(/Users/thomas/miniforge3/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so, 2): no suitable image found. Did find: /Users/thomas/miniforge3/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so: mach-o, but wrong architecture /Users/thomas/miniforge3/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so: mach-o, but wrong architecture

Any ideas what could be the issue here?

— You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub https://github.com/google/jax/issues/5501#issuecomment-885960599, or unsubscribe https://github.com/notifications/unsubscribe-auth/AACVGBB2TJRSBCDOATQAEJDTZH3L7ANCNFSM4WQWPQGA .

moellenh commented 3 years ago

Strange. I built it exactly as in dcxSt's post above and running the build script mentions "arm64" as a target. Perhaps there was some issue that I ran things accidentally in a Rosetta terminal, will double check this.

Update: I double checked, and everything was run in a regular terminal.

Update2: It runs now! :) For anyone having similar problems, my issue seems to have been that I installed bazel from a rosetta terminal. 

(base) thomas@mbpro jax % python3 build/build.py                                                                

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\

Starting local Bazel server and connecting to it...
Bazel binary path: /usr/local/bin/bazel
Python binary path: /Users/thomas/miniforge3/bin/python3
Python version: 3.9
NumPy version: 1.21.1
SciPy version: 1.7.0
MKL-DNN enabled: yes
Target CPU: arm64
Target CPU features: default
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
/usr/local/bin/bazel run --verbose_failures=true --config=short_logs --config=mkl_open_source_only :build_wheel -- --output_path=/Users/thomas/Documents/jax/dist --cpu=arm64

The resulting wheel is "jaxlib-0.1.70-cp39-none-macosx_11_0_arm64.whl", so it doesn't look like a x86 one.

gerdm commented 3 years ago

I tried to install it using @dcxSt's instructions, but I get the following error:

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\

Extracting Bazel installation...
Starting local Bazel server and connecting to it...
Bazel binary path: /usr/local/bin/bazel
Python binary path: /Users/gerardoduran/miniforge3/envs/pyprobml/bin/python
Python version: 3.9
NumPy version: 1.19.5
SciPy version: 1.7.0
MKL-DNN enabled: yes
Target CPU: arm64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
/usr/local/bin/bazel run --verbose_failures=true --config=short_logs --config=mkl_open_source_only :build_wheel -- --output_path=/Users/gerardoduran/Documents/repos/external/jax/dist --cpu=arm64
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /Users/gerardoduran/Documents/repos/external/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /Users/gerardoduran/Documents/repos/external/jax/.bazelrc:
  Inherited 'build' options: --repo_env PYTHON_BIN_PATH=/Users/gerardoduran/miniforge3/envs/pyprobml/bin/python --action_env=PYENV_ROOT --python_path=/Users/gerardoduran/miniforge3/envs/pyprobml/bin/python --repo_env TF_NEED_CUDA=0 --action_env TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 --repo_env TF_NEED_ROCM=0 --action_env TF_ROCM_AMDGPU_TARGETS=gfx803,gfx900,gfx906,gfx1010 -c opt --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true --spawn_strategy=standalone --strategy=Genrule=standalone --enable_platform_specific_config --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /Users/gerardoduran/Documents/repos/external/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:mkl_open_source_only in file /Users/gerardoduran/Documents/repos/external/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /Users/gerardoduran/Documents/repos/external/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /Users/gerardoduran/Documents/repos/external/jax/.bazelrc: --copt=-Wno-sign-compare --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading: 
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
WARNING: Download from http://mirror.tensorflow.org/github.com/tensorflow/runtime/archive/d29d1ef0a65a8f9c23e1f88067ce4205d3085e87.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found
Loading: 0 packages loaded
Loading: 0 packages loaded
Loading: 0 packages loaded
Analyzing: target //build:build_wheel (1 packages loaded, 0 targets configured)
DEBUG: Rule 'io_bazel_rules_docker' indicated that a canonical reproducible form can be obtained by modifying arguments shallow_since = "1596824487 -0400"
DEBUG: Repository io_bazel_rules_docker instantiated at:
  /Users/gerardoduran/Documents/repos/external/jax/WORKSPACE:34:10: in <toplevel>
  /private/var/tmp/_bazel_gerardoduran/fdbb951abd3afe67a47ab70535de95bf/external/org_tensorflow/tensorflow/workspace0.bzl:108:34: in workspace
  /private/var/tmp/_bazel_gerardoduran/fdbb951abd3afe67a47ab70535de95bf/external/bazel_toolchains/repositories/repositories.bzl:35:23: in repositories
Repository rule git_repository defined at:
  /private/var/tmp/_bazel_gerardoduran/fdbb951abd3afe67a47ab70535de95bf/external/bazel_tools/tools/build_defs/repo/git.bzl:199:33: in <toplevel>
Analyzing: target //build:build_wheel (14 packages loaded, 12 targets configured)
Analyzing: target //build:build_wheel (14 packages loaded, 12 targets configured)
Analyzing: target //build:build_wheel (14 packages loaded, 12 targets configured)
Analyzing: target //build:build_wheel (14 packages loaded, 12 targets configured)
Analyzing: target //build:build_wheel (47 packages loaded, 255 targets configured)
Analyzing: target //build:build_wheel (159 packages loaded, 4927 targets configured)
Analyzing: target //build:build_wheel (164 packages loaded, 6975 targets configured)
Analyzing: target //build:build_wheel (164 packages loaded, 6975 targets configured)
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/74da7ae0601728d7996e37c1f1828096e3d19103.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found
INFO: Analyzed target //build:build_wheel (172 packages loaded, 14115 targets configured).

INFO: Found 1 target...
[0 / 15] [Prepa] BazelWorkspaceStatusAction stable-status.txt ... (3 actions, 0 running)
ERROR: /private/var/tmp/_bazel_gerardoduran/fdbb951abd3afe67a47ab70535de95bf/external/com_google_protobuf/BUILD:301:11: Compiling src/google/protobuf/compiler/cpp/cpp_message.cc failed: (Exit 1): cc_wrapper.sh failed: error executing command 
  (cd /private/var/tmp/_bazel_gerardoduran/fdbb951abd3afe67a47ab70535de95bf/execroot/__main__ && \
  exec env - \
    PATH=/Users/gerardoduran/miniforge3/envs/pyprobml/bin:/Users/gerardoduran/miniforge3/condabin:/Users/gerardoduran/google-cloud-sdk/bin:/usr/local/opt/ruby/bin:/Users/gerardoduran/pomo:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin:/Library/TeX/texbin \
    PWD=/proc/self/cwd \
    TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 \
    TF_ROCM_AMDGPU_TARGETS=gfx803,gfx900,gfx906,gfx1010 \
  external/local_config_cc/cc_wrapper.sh -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections '-std=c++11' -MD -MF bazel-out/darwin-opt/bin/external/com_google_protobuf/_objs/protoc_lib/cpp_message.d '-frandom-seed=bazel-out/darwin-opt/bin/external/com_google_protobuf/_objs/protoc_lib/cpp_message.o' -iquote external/com_google_protobuf -iquote bazel-out/darwin-opt/bin/external/com_google_protobuf -iquote external/zlib -iquote bazel-out/darwin-opt/bin/external/zlib -isystem external/com_google_protobuf/src -isystem bazel-out/darwin-opt/bin/external/com_google_protobuf/src -isystem external/zlib -isystem bazel-out/darwin-opt/bin/external/zlib -Wno-sign-compare '-std=c++14' -DHAVE_PTHREAD -DHAVE_ZLIB -Woverloaded-virtual -Wno-sign-compare -Wno-unused-function -Wno-write-strings -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/com_google_protobuf/src/google/protobuf/compiler/cpp/cpp_message.cc -o bazel-out/darwin-opt/bin/external/com_google_protobuf/_objs/protoc_lib/cpp_message.o)
Execution platform: @local_execution_config_platform//:platform
external/com_google_protobuf/src/google/protobuf/compiler/cpp/cpp_message.cc:1695:39: error: no member named 'kInlinedType' in 'google::protobuf::internal::FieldMetadata'
      type = internal::FieldMetadata::kInlinedType;
             ~~~~~~~~~~~~~~~~~~~~~~~~~^
external/com_google_protobuf/src/google/protobuf/compiler/cpp/cpp_message.cc:2167:41: error: no member named 'TYPE_STRING_INLINED' in namespace 'google::protobuf::internal'; did you mean 'TYPE_STRING_CORD'?
            processing_type = internal::TYPE_STRING_INLINED;
                              ~~~~~~~~~~^~~~~~~~~~~~~~~~~~~
                                        TYPE_STRING_CORD
/usr/local/include/google/protobuf/generated_message_table_driven.h:72:3: note: 'TYPE_STRING_CORD' declared here
  TYPE_STRING_CORD = 19,
  ^
external/com_google_protobuf/src/google/protobuf/compiler/cpp/cpp_message.cc:2182:41: error: no member named 'TYPE_BYTES_INLINED' in namespace 'google::protobuf::internal'; did you mean 'TYPE_BYTES_CORD'?
            processing_type = internal::TYPE_BYTES_INLINED;
                              ~~~~~~~~~~^~~~~~~~~~~~~~~~~~
                                        TYPE_BYTES_CORD
/usr/local/include/google/protobuf/generated_message_table_driven.h:74:3: note: 'TYPE_BYTES_CORD' declared here
  TYPE_BYTES_CORD = 21,
  ^
3 errors generated.
Target //build:build_wheel failed to build
INFO: Elapsed time: 159.099s, Critical Path: 8.83s
INFO: 97 processes: 55 internal, 42 local.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
b''
Traceback (most recent call last):
  File "/Users/gerardoduran/Documents/repos/external/jax/build/build.py", line 604, in <module>
    main()
  File "/Users/gerardoduran/Documents/repos/external/jax/build/build.py", line 599, in main
    shell(command)
  File "/Users/gerardoduran/Documents/repos/external/jax/build/build.py", line 52, in shell
    output = subprocess.check_output(cmd)
  File "/Users/gerardoduran/miniforge3/envs/pyprobml/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/Users/gerardoduran/miniforge3/envs/pyprobml/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['/usr/local/bin/bazel', 'run', '--verbose_failures=true', '--config=short_logs', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/Users/gerardoduran/Documents/repos/external/jax/dist', '--cpu=arm64']' returned non-zero exit status 1.

I'm not sure if this error is because of the 404:

WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/74da7ae0601728d7996e37c1f1828096e3d19103.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found

Any idea why?

ahfs commented 3 years ago

@gerdm Missing Xcode could be the reason. If you haven't installed Xcode, get it from here. Make sure you open it to finish the installation. And then follow the steps in @dcxSt's comment. It worked for me.

koubadomik commented 3 years ago

@gerdm you should install xcode (not only cmd tools) - as mentioned https://github.com/google/jax/issues/5501#issuecomment-847237225. Following these steps worked for me.

egorssed commented 3 years ago

Hey, I acted as dcxSt advised. I had the same problem as gerdm. I managed to install everything without errors using the following sequence of actions. Installed Xcode like here installed homebrew and Bazel like here, than I acted according to dcxSt's guide.

The problem is that import jaxlib works, whereas import jax fails

 >>> import jax
/Users/egordanilov/Desktop/Science/Git/jax/jax/lib/__init__.py:31: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/egordanilov/Desktop/Science/Git/jax/jax/__init__.py", line 37, in <module>
    from . import config as _config_module
  File "/Users/egordanilov/Desktop/Science/Git/jax/jax/config.py", line 18, in <module>
    from jax._src.config import config
  File "/Users/egordanilov/Desktop/Science/Git/jax/jax/_src/config.py", line 27, in <module>
    from jax import lib
  File "/Users/egordanilov/Desktop/Science/Git/jax/jax/lib/__init__.py", line 70, in <module>
    from jaxlib import cpu_feature_guard
ImportError: dlopen(/usr/local/Caskroom/miniforge/base/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so, 2): no suitable image found.  Did find:
    /usr/local/Caskroom/miniforge/base/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so: mach-o, but wrong architecture
    /usr/local/Caskroom/miniforge/base/lib/python3.9/site-packages/jaxlib/cpu_feature_guard.so: mach-o, but wrong architecture

I tried to do pip install jax==0.2.10 advise from here, but had no success with it.

Does anyone have any solutions in mind ?

soerenab commented 3 years ago

Hey @egorssed, I ran into the same error as you. In contrast to you, I had actually initially not paid enough attention to how to install Xcode, Homebrew and Bazel. Once I followed those steps closely, the installation of jax worked out successfully. In particular, I believe that I had previously used a version of Bazel that was not meant for ARM64 / M1 architecture.

I also wrote down all steps in detail again, perhaps this makes it easier for others as well: (don't want to take false credits here though: all steps are taken from comments above or pages that are linked to in the comments above)

Edit: Apparently, these steps not always work flawlessly. Check out the comment below for small modifications of these instructions if you run into problems ;)

  1. Install xcode through the App Store. Once it is installed, open xcode and agree to the licence terms

  2. Install Rosetta2 emulator From what I understand Rosetta2 is a program by Apple that translates machine language that is written for Intel hardware (which I think it is called ‘x84-64’, but not 100% sure) into machine language that is appropriate for Apple’s M1 chip (which I think it is called ‘ARM64’ but also not 100% sure). You can install it via terminal with: % /usr/sbin/softwareupdate --install-rosetta --agree-to-license

  3. Install homebrew You will need to install two versions of homebrew:

    • Homebrew for the native M1 architecture. Install via terminal with: % arch -x86_64 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)" This version will (and must) be installed under /opt/homebrew .
    • Homebrew for the Intel architecture (from my understanding programs installed with this version of homebrew will then be translated by Rosetta2 once you execute them for the first time). Install via terminal with: % arch -x86_64 zsh % cd /usr/local && mkdir homebrew % curl -L https://github.com/Homebrew/brew/tarball/master | tar xz --strip 1 -C homebrew This version will (and must) be installed under /usr/local .

  4. Install python for arm64

    • If you already installed anaconda, you will have to remove it again. (Maybe there is a way to avoid this but this is what I did. The problem with anaconda is that it does not supply packages native for ARM64. We will instead install miniforge, which is like anaconda (from my understanding) but which in addition to the regular anaconda channels also supplies packages compiled for ARM64 hardware - in fact this is the default. See: https://stackoverflow.com/questions/60532678/what-is-the-difference-between-miniconda-and-miniforge)
    • Sanity check: % which brew should return % /opt/homebrew/bin/brew
    • Install miniforge via terminal with: % brew install --cask miniforge % conda init zsh % conda activate

  5. Create and activate new conda environment % conda create --name env_jax python=3.8 % conda activate env_jax

  6. Install numpy, scipy, and scikit-learn % conda install numpy scipy scikit-learn

  7. Install bazel % brew install bazel

  8. Compile and install jax from source % git clone https://github.com/google/jax % cd jax % python build/build.py % pip install dist/*.whl % pip install -e .

egorssed commented 3 years ago

@soerenab, many thanks to you! Painstakingly overcoming an urge to break my head against the wall I finally succeeded in installing Jax!

However, I still had to make some corrections to your installation algorithm.

3.Install homebrew I had some "Fetch error" with the command arch -x86_64 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install.sh)" I remember that I've seen that some related scripts have been deprecated, so I had to use this one arch -x86_64 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"

4.Install python for arm64

So there is also this usual problem with conda init zsh doing nothing and conda activate claiming that the shell is not configured.

The solution is to give an explicit source before handling Conda source /usr/local/Caskroom/miniforge/base/bin/activate

7.Install bazel There were also some problems with Bazel installation, but, apparently, it is not needed at all. When you build jax, it installs Bazel on its own

 >>> python3 build/build.py
...
Downloading bazel from: https://github.com/bazelbuild/bazel/releases/download/4.1.0/bazel-4.1.0-darwin-arm64
bazel-4.1.0-darwin-arm64 [########################################] 100%
...
ericmjl commented 3 years ago

Friends, just wanted to let you know that @xhochy has done this PR: https://github.com/conda-forge/jaxlib-feedstock/pull/60.

Please consider it community-supported and not officially-by-Google, which means also please go easy on @xhochy. The best thing you could do here is to offer to help in whatever way he deems necessary - that PR was a tour-de-force which was not easy at all. (I'm still studying the PR and don't fully understand everything that went in there myself!)

dfm commented 3 years ago

I spent some quality time with this thread getting jaxlib compiled on my new M1 mac but got it there eventually - thanks all!

Since I'm using jax for a few different projects, I decided to get the jaxlib v0.1.70 wheel building on GitHub Actions targeting arm64 (here's the workflow). I wanted to share this here in case it saves anyone time (vs. building from source) in the short term (hopefully there are official wheels soonish :D). It seems to be working for me when installed in a miniforge arm64 env, but I haven't put it through its paces too aggressively. You should be able to install it using:

python -m pip install "jax[cpu]" -f "https://dfm.io/custom-wheels/jaxlib/index.html"

Please feel free to open an issue on that repo (https://github.com/dfm/custom-wheels). Hope this is useful and sorry for the noise otherwise!

ericmjl commented 3 years ago

@dfm that's impressive! Can I ask, did you have some special trick for guaranteeing that you could build on arm64 hardware, or is that not necessary? I had always assumed that we'd need arm64 hardware to compile, and that Actions always gives us x86 hardware, but maybe I'm incorrect here!

dfm commented 3 years ago

@ericmjl: no, cross-compiling is supported so you should be able to build with a recent xcode on x64 and that's what I'm doing here. Testing remains an issue!

ericmjl commented 3 years ago

@dfm that's pretty cool. Thanks for educating me!

yashk2810 commented 3 years ago

JAX has now released a first version of jaxlib built for mac ARM architecture (v0.1.71).

(Preferred way) Upgrade your jax version to the latest version (pip install -U jax) and then pip install jax[cpu] -f https://storage.googleapis.com/jax-releases/jax_releases.html.

OR

You can install it via pip install jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html

Please try it out and let us know if there are any issues :)

Note: Make sure your python version is 3.9 because that's the only supported python version for this wheel.

hawkinsp commented 3 years ago

Note also: we don't actually test these wheels. See above. We're building them as a convenience to the community but we're relying on y'all to test them for the moment...

gerdm commented 3 years ago

Hi @yashk2810!

I just followed your installation instructions (the Preferred way), but it fails to run on my end

In [1]: import jax
/Users/gerardoduran/miniforge3/lib/python3.9/site-packages/jax/lib/__init__.py:31: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
  warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "

In [2]: jax.numpy.sqrt(2)
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
'cyclone' is not a recognized processor for this target (ignoring processor)
LLVM ERROR: 64-bit code requested on a subtarget that doesn't support it!
[1]    38628 abort      ipython

Do you know what might've caused this problem?

hawkinsp commented 3 years ago

@gerdm Nope, we'll have to debug that one. The way the wheels are built is by cross-compilation, as described above, so perhaps someone can see if they get the same behavior when building jaxlib 0.1.71 cross-compiled (or natively on arm64)? It appears it was working for 0.1.70, since folks reported that above.

We only have intermittent access to arm64 hardware, so we appreciate community help here!

yashk2810 commented 3 years ago

Looks like this problem was fixed in the rust repo like this: https://github.com/rust-lang/rust/pull/73086/files

hawkinsp commented 3 years ago

@yashk2810 Interesting. We just ask LLVM what the features of the local machine are: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc;drc=2876b0ebfd0080910d9eefb36d413cedb4a6ab8e;l=58 so it's not immediately obvious what we should do here.

dfm commented 3 years ago

I can confirm that I reproduce the same error as @gerdm when building from source on my M1. I haven't tried cross compiling yet. The same build environment worked for jaxlib v0.1.70 so it looks like an issue was introduced between versions. I'm not super familiar with the whole bazel build process, so I'm not sure how helpful I can be for tracking it down, but happy to help if I can!

hawkinsp commented 3 years ago

@dfm Well that's good news and bad news. It means our wheel build works (good!). But something obviously broke in JAX (really in XLA or LLVM, most likely... bad!)

yashk2810 commented 3 years ago

I built this wheel from TF head and jax head: https://storage.googleapis.com/jax-releases/mac/jaxlib-0.1.72-cp39-none-macosx_11_0_arm64.whl

Can you try this out and see if the error goes away?

Thank you!