jax: Build failures at HEAD on Apple ARM64

• If applicable, include full error messages/tracebacks.

Current top of main does not build on my machine (Apple M1 Pro, macOS 12.3.1):

jax on  main [⇡$!] via jax 
➜ python build/build.py                            

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


b'\x1b[31mERROR: The project you\'re trying to build requires Bazel 5.1.0 (specified in /Users/nicholasjunge/Workspaces/python/jax/.bazelversion), but it wasn\'t found in /opt/homebrew/Cellar/bazel/5.0.0/libexec/bin.\x1b[0m\n\nBazel binaries for all official releases can be downloaded from here:\n  https://github.com/bazelbuild/bazel/releases\n\nYou can download the required version directly using this command:\n  (cd "/opt/homebrew/Cellar/bazel/5.0.0/libexec/bin" && curl -fLO https://releases.bazel.build/5.1.0/release/bazel-5.1.0-darwin-arm64 && chmod +x bazel-5.1.0-darwin-arm64)\n'
Bazel binary path: ./bazel-5.1.0-darwin-arm64
Bazel version: 5.1.0
Python binary path: /Users/nicholasjunge/Workspaces/python/jax/venv/bin/python
Python version: 3.9
NumPy version: 1.22.3
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...
./bazel-5.1.0-darwin-arm64 run --verbose_failures=true --toolchain_resolution_debug=@bazel_tools//tools/cpp:toolchain_type --config=mkl_open_source_only :build_wheel -- --output_path=/Users/nicholasjunge/Workspaces/python/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/nicholasjunge/Workspaces/python/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /Users/nicholasjunge/Workspaces/python/jax/.bazelrc:
  Inherited 'build' options: --apple_platform_type=macos --macos_minimum_os=10.9 --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 -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
INFO: Reading rc options for 'run' from /Users/nicholasjunge/Workspaces/python/jax/.jax_configure.bazelrc:
  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/Users/nicholasjunge/Workspaces/python/jax/venv/bin/python --action_env=PYENV_ROOT --python_path=/Users/nicholasjunge/Workspaces/python/jax/venv/bin/python --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:mkl_open_source_only in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading: 
Loading: 0 packages loaded
INFO: Build option --distinct_host_configuration has changed, discarding analysis cache.
Analyzing: target //build:build_wheel (0 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/nicholasjunge/Workspaces/python/jax/WORKSPACE:37:14: in <toplevel>
  /private/var/tmp/_bazel_nicholasjunge/270a4a78734ae0f3124fa7265b8a65ef/external/org_tensorflow/tensorflow/workspace0.bzl:107:34: in workspace
  /private/var/tmp/_bazel_nicholasjunge/270a4a78734ae0f3124fa7265b8a65ef/external/bazel_toolchains/repositories/repositories.bzl:35:23: in repositories
Repository rule git_repository defined at:
  /private/var/tmp/_bazel_nicholasjunge/270a4a78734ae0f3124fa7265b8a65ef/external/bazel_tools/tools/build_defs/repo/git.bzl:199:33: in <toplevel>
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-armeabi-v7a; mismatching values: arm
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-armeabi-v7a; mismatching values: arm
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-darwin_arm64; mismatching values: arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-darwin_arm64; mismatching values: arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-darwin_arm64e; mismatching values: arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-darwin_arm64e; mismatching values: arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-darwin_x86_64; mismatching values: x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-darwin_x86_64; mismatching values: x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_arm64; mismatching values: ios, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_arm64; mismatching values: ios, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_arm64e; mismatching values: ios, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_arm64e; mismatching values: ios, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_armv7; mismatching values: ios, armv7
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_armv7; mismatching values: ios, armv7
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_i386; mismatching values: ios, i386
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_i386; mismatching values: ios, i386
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_sim_arm64; mismatching values: ios, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_sim_arm64; mismatching values: ios, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_x86_64; mismatching values: ios, x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-ios_x86_64; mismatching values: ios, x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-tvos_arm64; mismatching values: tvos, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-tvos_arm64; mismatching values: tvos, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-tvos_sim_arm64; mismatching values: tvos, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-tvos_sim_arm64; mismatching values: tvos, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-tvos_x86_64; mismatching values: tvos, x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-tvos_x86_64; mismatching values: tvos, x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_arm64; mismatching values: watchos, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_arm64; mismatching values: watchos, arm64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_arm64_32; mismatching values: watchos, arm64_32
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_arm64_32; mismatching values: watchos, arm64_32
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_armv7k; mismatching values: watchos, armv7k
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_armv7k; mismatching values: watchos, armv7k
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_i386; mismatching values: watchos, i386
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_i386; mismatching values: watchos, i386
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_x86_64; mismatching values: watchos, x86_64
INFO: ToolchainResolution:     Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: Rejected toolchain @local_config_cc//:cc-compiler-watchos_x86_64; mismatching values: watchos, x86_64
INFO: ToolchainResolution:   Type @bazel_tools//tools/cpp:toolchain_type: target platform @local_config_platform//:host: No toolchains found.
ERROR: /Users/nicholasjunge/Workspaces/python/jax/build/BUILD.bazel:25:10: While resolving toolchains for target //build:build_wheel: No matching toolchains found for types @bazel_tools//tools/cpp:toolchain_type. Maybe --incompatible_use_cc_configure_from_rules_cc has been flipped and there is no default C++ toolchain added in the WORKSPACE file? See https://github.com/bazelbuild/bazel/issues/10134 for details and migration instructions.
ERROR: Analysis of target '//build:build_wheel' failed; build aborted: 
INFO: Elapsed time: 0.093s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded, 253 targets configured)
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully (0 packages loaded, 253 targets configured)
b''
Traceback (most recent call last):
  File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 527, in <module>
    main()
  File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 522, in main
    shell(command)
  File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/opt/homebrew/Cellar/python@3.9/3.9.12/Frameworks/Python.framework/Versions/3.9/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/opt/homebrew/Cellar/python@3.9/3.9.12/Frameworks/Python.framework/Versions/3.9/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-5.1.0-darwin-arm64', 'run', '--verbose_failures=true', '--toolchain_resolution_debug=@bazel_tools//tools/cpp:toolchain_type', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/Users/nicholasjunge/Workspaces/python/jax/dist', '--cpu=arm64']' returned non-zero exit status 1.

Some notes / observations:

• Apparently, toolchain resolution does not work properly for downloaded versions of Bazel; currently, the latest Homebrew-supplied Bazel is still v5.0.0.
• Editing the .bazelversion (and, as of recently, build.py) to use Bazel 5.0.0 also fails, with a different error (it complains about a win_ver_file setting in Tensorflow’s BUILD file, which I guess was not supported with v5.0.0).
• The above output was generated by injecting the --toolchain_resolution_debug=@bazel_tools//tools/cpp:toolchain_type flag into the build args. Somehow, all candidate toolchains are rejected, but I don’t know why.
• I switched over to a local Tensorflow repo on my machine (as mentioned in the WORKSPACE file). Apparently in that setting, Bazel uses whatever ref is checked out there (currently, that is master for me, at or near current HEAD). So I pulled in some of the recent changes in Tensorflow’s build config.

➜ git log
commit e1bbbf55cd28abe60175a5cf106c4c2c9c9e044f (HEAD -> main, upstream/main)
Merge: 359b614b c72d8f6b
Author: jax authors <no-reply@google.com>
Date:   Sun Apr 3 12:11:03 2022 -0700
 
    Merge pull request #10130 from mattjj:no-string-annotations
    
    PiperOrigin-RevId: 439174012