Issue 1:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

python build/build.py

Building jaxlib at HEAD on my machine using bazel 4.1.0 fails (macOS 12.0, Macbook Pro 14" M1 Pro 2021). This error seems to come from a bad xcode-locator call, possibly because the bazel version that is pulled there (see the traceback below) is configured for the wrong architecture? IIRC, I also do not (yet) have Rosetta installed.

Manually editing the .bazelversion file and setting it to 4.2.1 (which enables the use of the latest bazel version provided by brew) avoids this issue and builds the wheel successfully (though I have not tested it yet).

If applicable, include full error messages/tracebacks.

➜ python build/build.py

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

b'\x1b[31mERROR: The project you\'re trying to build requires Bazel 4.1.0 (specified in /Users/nicholasjunge/Workspaces/python/jax/.bazelversion), but it wasn\'t found in /opt/homebrew/Cellar/bazel/4.2.1_1/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/4.2.1_1/libexec/bin" && curl -fLO https://releases.bazel.build/4.1.0/release/bazel-4.1.0-darwin-arm64 && chmod +x bazel-4.1.0-darwin-arm64)\n'
Bazel binary path: ./bazel-4.1.0-darwin-arm64
Bazel version: 4.1.0
Python binary path: /Users/nicholasjunge/Workspaces/python/jax/venv/bin/python
Python version: 3.9
NumPy version: 1.21.4
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-4.1.0-darwin-arm64 run --verbose_failures=true --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 --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
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/8909dc5ebe8ad39f1743131eb70df402d796acab.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/tensorflow/runtime/archive/fbdd15997f4495011fed44d051152a2a69436e69.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found
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:108: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>
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/open-source-parsers/jsoncpp/archive/1.9.5.tar.gz failed: class com.google.devtools.build.lib.bazel.repository.downloader.UnrecoverableHttpException GET returned 404 Not Found
INFO: Analyzed target //build:build_wheel (0 packages loaded, 0 targets configured).
INFO: Found 1 target...
[0 / 5] [Prepa] BazelWorkspaceStatusAction stable-status.txt ... (2 actions, 0 running)
ERROR: /private/var/tmp/_bazel_nicholasjunge/270a4a78734ae0f3124fa7265b8a65ef/external/zlib/BUILD.bazel:5:11: Compiling gzclose.c failed: Exec failed due to IOException: com.google.devtools.build.lib.shell.ExecFailedException: java.io.IOException: Cannot run program "/var/tmp/_bazel_nicholasjunge/install/d0b631b3e5a643567f221e133629d965/xcode-locator": error=86, Bad CPU type in executable
Target //build:build_wheel failed to build
INFO: Elapsed time: 0.249s, Critical Path: 0.04s
INFO: 20 processes: 19 internal, 1 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/nicholasjunge/Workspaces/python/jax/build/build.py", line 524, in <module>
main()
File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 519, 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.8/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.8/Frameworks/Python.framework/Versions/3.9/lib/python3.9/subprocess.py", line 528, in run
raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-4.1.0-darwin-arm64', 'run', '--verbose_failures=true', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/Users/nicholasjunge/Workspaces/python/jax/dist', '--cpu=arm64']' returned non-zero exit status 36.

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

import jax.numpy as jnp
jnp.flip(jnp.zeros((3, 3, 3)), axis=(1,2))

If applicable, include full error messages/tracebacks.

Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/home/buehrle/dev/polyssimo/.venv/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py", line 1114, in flip
return lax.rev(m, [_canonicalize_axis(axis, ndim(m))])
File "/home/buehrle/dev/polyssimo/.venv/lib/python3.6/site-packages/jax/_src/util.py", line 262, in canonicalize_axis
axis = operator.index(axis)
TypeError: 'tuple' object cannot be interpreted as an integer