I am trying to install jax on a Mac M2. pip install jax works, but results in an error that states:
RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able to work around this issue by building Jaxlib from source.
I tried following the instructions at Accelerated JAX training on Mac, but it failed with
Bazel binary path: /usr/local/bin/bazel
Bazel version: 5.1.1
Python binary path: /Users/jkitchin/opt/anaconda3/bin/python
Python version: 3.9
NumPy version: 1.24.4
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no
Plugin device enabled: no
Building XLA and installing it in the jaxlib source tree...
/usr/local/bin/bazel run --verbose_failures=true :build_wheel -- --output_path=/Users/jkitchin/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/jkitchin/jax/.bazelrc:
Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /Users/jkitchin/jax/.bazelrc:
Inherited 'build' options: --nocheck_visibility --apple_platform_type=macos --macos_minimum_os=10.14 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --experimental_cc_shared_library --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true --define=tsl_link_protobuf=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir. --@xla//xla/python:enable_gpu=false --@xla//xla/python:enable_tpu=false --@xla//xla/python:enable_plugin_device=false
INFO: Reading rc options for 'run' from /Users/jkitchin/jax/.jax_configure.bazelrc:
Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/Users/jkitchin/opt/anaconda3/bin/python --action_env=PYENV_ROOT --python_path=/Users/jkitchin/opt/anaconda3/bin/python --distinct_host_configuration=false --@xla//xla/python:enable_tpu=true --config=avx_posix --config=mkl_open_source_only
INFO: Found applicable config definition build:short_logs in file /Users/jkitchin/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:avx_posix in file /Users/jkitchin/jax/.bazelrc: --copt=-mavx --host_copt=-mavx
INFO: Found applicable config definition build:mkl_open_source_only in file /Users/jkitchin/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /Users/jkitchin/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /Users/jkitchin/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++17 --host_cxxopt=-std=c++17
DEBUG: /private/var/tmp/_bazel_jkitchin/e6abbba145bbf6e2a60f41057ec994c8/external/xla/third_party/repo.bzl:132:14:
Warning: skipping import of repository 'tf_runtime' because it already exists.
DEBUG: /private/var/tmp/_bazel_jkitchin/e6abbba145bbf6e2a60f41057ec994c8/external/xla/third_party/repo.bzl:132:14:
Warning: skipping import of repository 'llvm-raw' because it already exists.
Loading:
Loading: 0 packages loaded
Analyzing: target //build:build_wheel (1 packages loaded, 0 targets configured)
Analyzing: target //build:build_wheel (7 packages loaded, 199 targets configured)
INFO: Analyzed target //build:build_wheel (8 packages loaded, 323 targets configured).
INFO: Found 1 target...
[0 / 1] [Prepa] BazelWorkspaceStatusAction stable-status.txt
ERROR: /private/var/tmp/_bazel_jkitchin/e6abbba145bbf6e2a60f41057ec994c8/external/xla/xla/BUILD:83:17: ProtoCompile external/xla/xla/xla_data.pb.h failed: (Illegal instruction): protoc failed: error executing command
(cd /private/var/tmp/_bazel_jkitchin/e6abbba145bbf6e2a60f41057ec994c8/execroot/__main__ && \
exec env - \
PATH=/opt/homebrew/Cellar/pyenv/2.3.21/libexec:/opt/homebrew/Cellar/pyenv/2.3.21/plugins/python-build/bin:/opt/homebrew/Cellar/pyenv-virtualenv/1.2.1/shims:/Users/jkitchin/.pyenv/shims:/Users/jkitchin/opt/anaconda3/bin:/Users/jkitchin/opt/anaconda3/condabin:/opt/homebrew/bin:/opt/homebrew/sbin:/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin \
PYENV_ROOT=/Users/jkitchin/.pyenv \
bazel-out/darwin-opt-exec-50AE0418/bin/external/com_google_protobuf/protoc '--cpp_out=bazel-out/darwin-opt/bin/external/xla' -Iexternal/xla -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/any_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/source_context_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/type_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/api_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/descriptor_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/compiler_plugin_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/duration_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/empty_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/field_mask_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/struct_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/timestamp_proto -Ibazel-out/darwin-opt/bin/external/com_google_protobuf/_virtual_imports/wrappers_proto external/xla/xla/xla_data.proto)
# Configuration: 3b261c317d2ca1663cdefd6c23d3ada0f18e752c22e20ab2112a26d4805801da
# Execution platform: @local_execution_config_platform//:platform
Target //build:build_wheel failed to build
INFO: Elapsed time: 4.659s, Critical Path: 0.09s
INFO: 12 processes: 12 internal.
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/jkitchin/jax/build/build.py", line 564, in <module>
main()
File "/Users/jkitchin/jax/build/build.py", line 559, in main
shell(command)
File "/Users/jkitchin/jax/build/build.py", line 53, in shell
output = subprocess.check_output(cmd)
File "/Users/jkitchin/opt/anaconda3/lib/python3.9/subprocess.py", line 424, in check_output
return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
File "/Users/jkitchin/opt/anaconda3/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', ':build_wheel', '--', '--output_path=/Users/jkitchin/jax/dist', '--cpu=x86_64']' returned non-zero exit status 1.
I guess '--cpu=x86_64' isn't right, since that is not the CPU on the machine, but somewhere it seems to be the default.
