Issue 1:
Please see the following error messages regarding the error during the build. The error disappears and the build is successful after manually fixing the BUILD file.

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 --enable_cuda

If applicable, include full error messages/tracebacks.

(build_jax) /export/share/software/jax> python build/build.py --enable_cuda

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

Bazel binary path: /workspace/anaconda3/envs/build_jax/bin/bazel
Bazel version: 5.0.0
Python binary path: /workspace/anaconda3/envs/build_jax/bin/python
Python version: 3.8
NumPy version: 1.20.3
MKL-DNN enabled: yes
Target CPU: ppc64le
Target CPU features: release
CUDA enabled: yes
CUDA compute capabilities: 3.5,5.2,6.0,7.0,8.0
NCCL enabled: yes
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
/workspace/anaconda3/envs/build_jax/bin/bazel run --verbose_failures=true --config=mkl_open_source_only --config=cuda :build_wheel -- --output_path=/export/share/software/jax/dist --cpu=ppc64le
Starting local Bazel server and connecting to it...
INFO: Options provided by the client:
Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /export/share/software/jax/.bazelrc:
Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /export/share/software/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 /export/share/software/jax/.jax_configure.bazelrc:
Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/workspace/anaconda3/envs/build_jax/bin/python --action_env=PYENV_ROOT --python_path=/workspace/anaconda3/envs/build_jax/bin/python --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /export/share/software/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:mkl_open_source_only in file /export/share/software/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:cuda in file /export/share/software/jax/.bazelrc: --repo_env TF_NEED_CUDA=1 --action_env TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 --crosstool_top=@local_config_cuda//crosstool:toolchain --@local_config_cuda//:enable_cuda --define=xla_python_enable_gpu=true
INFO: Found applicable config definition build:cuda in file /export/share/software/jax/.jax_configure.bazelrc: --action_env TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,7.0,8.0
INFO: Found applicable config definition build:linux in file /export/share/software/jax/.bazelrc: --config=posix --copt=-Wno-stringop-truncation --copt=-Wno-array-parameter
INFO: Found applicable config definition build:posix in file /export/share/software/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
WARNING: Download from https://storage.googleapis.com/mirror.tensorflow.org/github.com/tensorflow/runtime/archive/102f5befef27a33dec2cab2600f9b2b4bf3e5036.tar.gz failed: class java.io.FileNotFoundException GET returned 404 Not Found
WARNING: Download from https://mirror.bazel.build/github.com/bazelbuild/rules_cc/archive/081771d4a0e9d7d3aa0eed2ef389fa4700dfb23e.tar.gz failed: class java.io.FileNotFoundException GET returned 404 Not Found
DEBUG: /workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_tools/tools/cpp/lib_cc_configure.bzl:118:10:
Auto-Configuration Warning: 'TMP' environment variable is not set, using 'C:\Windows\Temp' as default
Loading:
Loading: 1 packages loaded
Analyzing: target //build:build_wheel (2 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:
/export/share/software/jax/WORKSPACE:37:14: in <toplevel>
/workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/org_tensorflow/tensorflow/workspace0.bzl:108:34: in workspace
/workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_toolchains/repositories/repositories.bzl:35:23: in repositories
Repository rule git_repository defined at:
/workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_tools/tools/build_defs/repo/git.bzl:199:33: in <toplevel>
ERROR: /workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_tools/tools/zip/BUILD:11:1: indentation error
ERROR: /workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_tools/tools/zip/BUILD:14:2: Trailing comma is allowed only in parenthesized tuples.
ERROR: /workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_tools/tools/zip/BUILD:14:3: syntax error at 'outdent': expected expression
ERROR: /export/share/software/jax/build/BUILD.bazel:25:10: every rule of type py_binary implicitly depends upon the target '@bazel_tools//tools/zip:zipper', but this target could not be found because of: no such target '@bazel_tools//tools/zip:zipper': target 'zipper' not declared in package 'tools/zip' defined by /workspace/anaconda3/envs/build_jax/share/bazel/66cc389786f6f90c4be41c4758c2d3a1/external/bazel_tools/tools/zip/BUILD
ERROR: Analysis of target '//build:build_wheel' failed; build aborted: Analysis failed
INFO: Elapsed time: 212.845s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (52 packages loaded, 212 targets configured)
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully (52 packages loaded, 212 targets configured)
b''
Traceback (most recent call last):
File "build/build.py", line 524, in <module>
main()
File "build/build.py", line 519, in main
shell(command)
File "build/build.py", line 53, in shell
output = subprocess.check_output(cmd)
File "/workspace/anaconda3/envs/build_jax/lib/python3.8/subprocess.py", line 415, in check_output
return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
File "/workspace/anaconda3/envs/build_jax/lib/python3.8/subprocess.py", line 516, in run
raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['/workspace/anaconda3/envs/build_jax/bin/bazel', 'run', '--verbose_failures=true', '--config=mkl_open_source_only', '--config=cuda', ':build_wheel', '--', '--output_path=/export/share/software/jax/dist', '--cpu=ppc64le']' returned non-zero exit status 1.

Issue 2:
Description

The code below computes an element on the diagonal of a pairwise distance matrix, with a modified sqrt function that has finite gradients at 0. Without JIT, the gradients wrt inputs are 0., as expected. With JIT, the gradients are ~1e12.

This shows up in TFP-on-JAX's GaussianProcess. As a workaround, we can set the diagonal of the distance matrix to 0 (which zeros out the gradients), but that doesn't account for the case where the input array has duplicate points (so zeros can appear elsewhere in the distance matrix). This is blocking Vizier's use of TFP-on-JAX's GP.

import jax
from jax import numpy as jnp
import numpy as np

x = jnp.array(
[-1.152883 , -1.229587 , -0.7161828, 1.6050261, -0.9966271, 1.0175147])

@jax.custom_gradient
def sqrt_with_finite_grads(x):
"""Avoid NaN gradients when `x == 0`."""
def grad(grad_ys):
large_float_like_x = np.sqrt(np.finfo(x.dtype).max)

# In this repro, x == 0, so safe_grads = large_float_like_x.
safe_grads = jnp.where(
jnp.equal(x, 0.), large_float_like_x, 0.5 / jnp.sqrt(x))
return grad_ys * safe_grads
return jnp.sqrt(x), grad

def dist(x):
"""Stripped down from https://github.com/tensorflow/probability/blob/main/tensorflow_probability/python/math/psd_kernels/internal/util.py#L191"""
# These two terms are equal. If we compute them differently (`sum(x*x)`, `dot(x, x)`, `sum(x**2)`) the issue persists. If we compute them the same way, the issue disappears.
# (This is how we compute the diagonal of a squared distance matrix, when both inputs are the same.)
row_norm_x = 2. * jnp.sum(x**2, axis=-1)
middle_term = 2. * jnp.dot(x, x)

# Their difference is 0.
out = row_norm_x - middle_term
d = jnp.clip(out, 0., jnp.inf)

# Clipping at a very small value, instead of 0., appears to fix it. (Because `clip` zeros out the gradient when the value is outside the interval?)
# d = jnp.clip(out, 1e-30, jnp.inf)

return sqrt_with_finite_grads(d)

print('No jit grads: ', jax.grad(dist)(x)) # => [0., 0., ...]
print('Jit grads: ', jax.jit(jax.grad(dist))(x)) # => [-1.9e12, -1.7e12, ...]

Colab with the same repro.

What jax/jaxlib version are you using?

jax v0.3.17

Which accelerator(s) are you using?

CPU

Additional System Info

py3.7