https://github.com/google/jax
Raw File
Tip revision: d91fe4caec6be197011fdd0a22cca5ba26d753a6 authored by Lena Martens on 25 January 2022, 17:23:30 UTC
Strip named_shape and weak_type from aval when donating buffers.
Tip revision: d91fe4c
.bazelrc
############################################################################
# All default build options below.

# Sets the default Apple platform to macOS.
build --apple_platform_type=macos
build --macos_minimum_os=10.9

# Make Bazel print out all options from rc files.
build --announce_rc

build --define open_source_build=true

build --spawn_strategy=standalone

build --enable_platform_specific_config

# Disable enabled-by-default TensorFlow features that we don't care about.
build --define=no_aws_support=true
build --define=no_gcp_support=true
build --define=no_hdfs_support=true
build --define=no_kafka_support=true
build --define=no_ignite_support=true

build --define=grpc_no_ares=true

build -c opt

build --config=short_logs

build --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.

###########################################################################

build:posix --copt=-fvisibility=hidden
build:posix --copt=-Wno-sign-compare
build:posix --cxxopt=-std=c++14
build:posix --host_cxxopt=-std=c++14

build:avx_posix --copt=-mavx
build:avx_posix --host_copt=-mavx

build:avx_windows --copt=/arch=AVX

build:avx_linux --copt=-mavx
build:avx_linux --host_copt=-mavx

build:native_arch_posix --copt=-march=native
build:native_arch_posix --host_copt=-march=native

build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1

build:cuda --repo_env TF_NEED_CUDA=1
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="3.5,5.2,6.0,6.1,7.0"
build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
build:cuda --@local_config_cuda//:enable_cuda
build:cuda --define=xla_python_enable_gpu=true

build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain
build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true
build:rocm --define=xla_python_enable_gpu=true
build:rocm --repo_env TF_NEED_ROCM=1
build:rocm --action_env TF_ROCM_AMDGPU_TARGETS="gfx900,gfx906,gfx908"

build:nonccl --define=no_nccl_support=true

# Tensorflow uses M_* math constants that only get defined by MSVC headers if
# _USE_MATH_DEFINES is defined.
build:windows --copt=/D_USE_MATH_DEFINES
build:windows --host_copt=/D_USE_MATH_DEFINES
# Make sure to include as little of windows.h as possible
build:windows --copt=-DWIN32_LEAN_AND_MEAN
build:windows --host_copt=-DWIN32_LEAN_AND_MEAN
build:windows --copt=-DNOGDI
build:windows --host_copt=-DNOGDI
# https://devblogs.microsoft.com/cppblog/announcing-full-support-for-a-c-c-conformant-preprocessor-in-msvc/
# otherwise, there will be some compiling error due to preprocessing.
build:windows --copt=/Zc:preprocessor
build:windows --cxxopt=/std:c++14
build:windows --host_cxxopt=/std:c++14
# Generate PDB files, to generate useful PDBs, in opt compilation_mode
# --copt /Z7 is needed.
build:windows --linkopt=/DEBUG
build:windows --host_linkopt=/DEBUG
build:windows --linkopt=/OPT:REF
build:windows --host_linkopt=/OPT:REF
build:windows --linkopt=/OPT:ICF
build:windows --host_linkopt=/OPT:ICF
build:windows --experimental_strict_action_env=true

build:linux --config=posix
# Workaround for gcc 10+ warnings related to upb.
# See https://github.com/tensorflow/tensorflow/issues/39467
build:linux --copt=-Wno-stringop-truncation
build:linux --copt=-Wno-array-parameter

build:macos --config=posix

# Suppress all warning messages.
build:short_logs --output_filter=DONT_MATCH_ANYTHING

build:tpu --define=with_tpu_support=true

#########################################################################
# RBE config options below.
# Flag to enable remote config
common --experimental_repo_remote_exec

build:rbe --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
build:rbe --google_default_credentials
build:rbe --bes_backend=buildeventservice.googleapis.com
build:rbe --bes_results_url="https://source.cloud.google.com/results/invocations"
build:rbe --bes_timeout=600s
build:rbe --define=EXECUTOR=remote
build:rbe --distinct_host_configuration=false
build:rbe --flaky_test_attempts=3
build:rbe --jobs=200
build:rbe --remote_executor=grpcs://remotebuildexecution.googleapis.com
build:rbe --remote_timeout=3600
build:rbe --spawn_strategy=remote,worker,standalone,local
test:rbe --test_env=USER=anon
# Attempt to minimize the amount of data transfer between bazel and the remote
# workers:
build:rbe --remote_download_toplevel

build:rbe_linux --config=rbe
build:rbe_linux --action_env=PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/go/bin"
build:rbe_linux --host_javabase=@bazel_toolchains//configs/ubuntu16_04_clang/1.1:jdk8
build:rbe_linux --javabase=@bazel_toolchains//configs/ubuntu16_04_clang/1.1:jdk8
build:rbe_linux --host_java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8
build:rbe_linux --java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8

# Non-rbe settings we should include because we do not run configure
build:rbe_linux --config=avx_linux
build:rbe_linux --linkopt=-lrt
build:rbe_linux --host_linkopt=-lrt
build:rbe_linux --linkopt=-lm
build:rbe_linux --host_linkopt=-lm

# Use the GPU toolchain until the CPU one is ready.
# https://github.com/bazelbuild/bazel/issues/13623
build:rbe_cpu_linux_base --config=rbe_linux
build:rbe_cpu_linux_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_cpu_linux_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_cpu_linux_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
build:rbe_cpu_linux_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
build:rbe_cpu_linux_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
build:rbe_cpu_linux_base --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"

build:rbe_cpu_linux_py37 --config=rbe_cpu_linux_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.7"
build:rbe_cpu_linux_py37 --python_path="/usr/local/bin/python3.7"
build:rbe_cpu_linux_py38 --config=rbe_cpu_linux_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.8"
build:rbe_cpu_linux_py38 --python_path="/usr/local/bin/python3.8"
build:rbe_cpu_linux_py39 --config=rbe_cpu_linux_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.9"
build:rbe_cpu_linux_py39 --python_path="/usr/local/bin/python3.9"
build:rbe_cpu_linux_py310 --config=rbe_cpu_linux_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.10"
build:rbe_cpu_linux_py310 --python_path="/usr/local/bin/python3.10"

build:rbe_linux_cuda_base --config=rbe_linux
build:rbe_linux_cuda_base --config=cuda
build:rbe_linux_cuda_base --repo_env=REMOTE_GPU_TESTING=1

build:rbe_linux_cuda11.1_nvcc_base --config=rbe_linux_cuda_base
build:rbe_linux_cuda11.1_nvcc_base --action_env=TF_CUDA_VERSION=11
build:rbe_linux_cuda11.1_nvcc_base --action_env=TF_CUDNN_VERSION=8
build:rbe_linux_cuda11.1_nvcc_base --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-11.1"
build:rbe_linux_cuda11.1_nvcc_base --action_env=LD_LIBRARY_PATH="/usr/local/cuda:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/tensorrt/lib"
build:rbe_linux_cuda11.1_nvcc_base --action_env=GCC_HOST_COMPILER_PATH="/dt7/usr/bin/gcc"
test:rbe_linux_cuda11.1_nvcc_base --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda-11.1/lib64"
build:rbe_linux_cuda11.1_nvcc_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_linux_cuda11.1_nvcc_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_linux_cuda11.1_nvcc_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
build:rbe_linux_cuda11.1_nvcc_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.1_nvcc_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.1_nvcc_base --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.1_nvcc_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_cuda"
build:rbe_linux_cuda11.1_nvcc_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_tensorrt"
build:rbe_linux_cuda11.1_nvcc_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_nccl"
build:rbe_linux_cuda11.1_nvcc_py3.7 --config=rbe_linux_cuda11.1_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_python3.7"
build:rbe_linux_cuda11.1_nvcc_py3.7 --python_path="/usr/local/bin/python3.7"
build:rbe_linux_cuda11.1_nvcc_py3.8 --config=rbe_linux_cuda11.1_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_python3.8"
build:rbe_linux_cuda11.1_nvcc_py3.8 --python_path="/usr/local/bin/python3.8"
build:rbe_linux_cuda11.1_nvcc_py3.9 --config=rbe_linux_cuda11.1_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.1-cudnn8-tensorrt7.2_config_python3.9"
build:rbe_linux_cuda11.1_nvcc_py3.9 --python_path="/usr/local/bin/python3.9"

build:rbe_linux_cuda11.4_nvcc_base --config=rbe_linux_cuda_base
build:rbe_linux_cuda11.4_nvcc_base --action_env=TF_CUDA_VERSION=11
build:rbe_linux_cuda11.4_nvcc_base --action_env=TF_CUDNN_VERSION=8
build:rbe_linux_cuda11.4_nvcc_base --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-11.4"
build:rbe_linux_cuda11.4_nvcc_base --action_env=LD_LIBRARY_PATH="/usr/local/cuda:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/tensorrt/lib"
build:rbe_linux_cuda11.4_nvcc_base --action_env=GCC_HOST_COMPILER_PATH="/dt7/usr/bin/gcc"
build:rbe_linux_cuda11.4_nvcc_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_linux_cuda11.4_nvcc_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_linux_cuda11.4_nvcc_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
build:rbe_linux_cuda11.4_nvcc_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_nvcc_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_nvcc_base --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_nvcc_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_cuda"
build:rbe_linux_cuda11.4_nvcc_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_tensorrt"
build:rbe_linux_cuda11.4_nvcc_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_nccl"
build:rbe_linux_cuda11.4_nvcc_py3.7 --config=rbe_linux_cuda11.4_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_python3.7"
build:rbe_linux_cuda11.4_nvcc_py3.7 --python_path="/usr/local/bin/python3.7"
build:rbe_linux_cuda11.4_nvcc_py3.8 --config=rbe_linux_cuda11.4_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_python3.8"
build:rbe_linux_cuda11.4_nvcc_py3.8 --python_path="/usr/local/bin/python3.8"
build:rbe_linux_cuda11.4_nvcc_py3.9 --config=rbe_linux_cuda11.4_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_python3.9"
build:rbe_linux_cuda11.4_nvcc_py3.9 --python_path="/usr/local/bin/python3.9"
build:rbe_linux_cuda11.4_nvcc_py3.10 --config=rbe_linux_cuda11.4_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.2-tensorrt7.2_config_python3.10"
build:rbe_linux_cuda11.4_nvcc_py3.10 --python_path="/usr/local/bin/python3.10"

build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --config=rbe_linux_cuda_base
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --action_env=TF_CUDA_VERSION=11
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --action_env=TF_CUDNN_VERSION=8.0.5
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-11.4"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --action_env=LD_LIBRARY_PATH="/usr/local/cuda:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/tensorrt/lib"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --action_env=GCC_HOST_COMPILER_PATH="/dt7/usr/bin/gcc"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_cuda//crosstool:toolchain"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_cuda"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_nccl"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.7 --config=rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_python3.7"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.7 --python_path="/usr/local/bin/python3.7"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.8 --config=rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_python3.8"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.8 --python_path="/usr/local/bin/python3.8"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.9 --config=rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_python3.9"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.9 --python_path="/usr/local/bin/python3.9"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.10 --config=rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base  --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_python3.10"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.10 --python_path="/usr/local/bin/python3.10"

# These you may need to change for your own GCP project.
build:tensorflow_testing_rbe --project_id=tensorflow-testing
common:tensorflow_testing_rbe_linux --remote_instance_name=projects/tensorflow-testing/instances/default_instance
build:tensorflow_testing_rbe_linux --config=tensorflow_testing_rbe
#############################################################################

# Load `.jax_configure.bazelrc` file written by build.py
try-import %workspace%/.jax_configure.bazelrc
back to top