removeRulesCC = false; <del> GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin"; GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc"; </del> <ins> GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${backend_cc_joined}/bin"; GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${backend_cc_joined}/bin/gcc"; </ins> # The version is automatically set to ".dev" if this variable is not set. # https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3 JAXLIB_RELEASE = "1"; <del> preConfigure = '' # dummy ldconfig mkdir dummy-ldconfig echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig chmod +x dummy-ldconfig/ldconfig export PATH="$PWD/dummy-ldconfig:$PATH" cat <<CFG > ./.jax_configure.bazelrc build --strategy=Genrule=standalone build --repo_env PYTHON_BIN_PATH="${python}/bin/python" build --action_env=PYENV_ROOT build --python_path="${python}/bin/python" build --distinct_host_configuration=false build --define PROTOBUF_INCLUDE_PATH="${pkgs.protobuf}/include" '' + lib.optionalString (stdenv.hostPlatform.avxSupport && stdenv.hostPlatform.isUnix) '' build --config=avx_posix '' + lib.optionalString mklSupport '' build --config=mkl_open_source_only '' + lib.optionalString cudaSupport '' build --action_env CUDA_TOOLKIT_PATH="${cudatoolkit_joined}" build --action_env CUDNN_INSTALL_PATH="${cudnn}" build --action_env TF_CUDA_PATHS="${cudatoolkit_joined},${cudnn},${nccl}" build --action_env TF_CUDA_VERSION="${lib.versions.majorMinor cudatoolkit.version}" build --action_env TF_CUDNN_VERSION="${lib.versions.major cudnn.version}" build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="${builtins.concatStringsSep "," cudaFlags.realArches}" '' + '' CFG ''; </del> <ins> preConfigure = # Dummy ldconfig to work around "Can't open cache file /nix/store/<hash>-glibc-2.38-44/etc/ld.so.cache" error '' mkdir dummy-ldconfig echo "#!${effectiveStdenv.shell}" > dummy-ldconfig/ldconfig chmod +x dummy-ldconfig/ldconfig export PATH="$PWD/dummy-ldconfig:$PATH" '' + # Construct .jax_configure.bazelrc. See https://github.com/google/jax/blob/b9824d7de3cb30f1df738cc42e486db3e9d915ff/build/build.py#L259-L345 # for more info. We assume # * `cpu = None` # * `enable_nccl = True` # * `target_cpu_features = "release"` # * `rocm_amdgpu_targets = None` # * `enable_rocm = False` # * `build_gpu_plugin = False` # * `use_clang = False` (Should we use `effectiveStdenv.cc.isClang` instead?) # # Note: We should try just running https://github.com/google/jax/blob/ceb198582b62b9e6f6bdf20ab74839b0cf1db16e/build/build.py#L259-L266 # instead of duplicating the logic here. Perhaps we can leverage the # `--configure_only` flag (https://github.com/google/jax/blob/ceb198582b62b9e6f6bdf20ab74839b0cf1db16e/build/build.py#L544-L548)? '' cat <<CFG > ./.jax_configure.bazelrc build --strategy=Genrule=standalone build --repo_env PYTHON_BIN_PATH="${python}/bin/python" build --action_env=PYENV_ROOT build --python_path="${python}/bin/python" build --distinct_host_configuration=false build --define PROTOBUF_INCLUDE_PATH="${pkgs.protobuf}/include" '' + lib.optionalString cudaSupport '' build --config=cuda build --action_env CUDA_TOOLKIT_PATH="${cuda_build_deps_joined}" build --action_env CUDNN_INSTALL_PATH="${cudnn}" build --action_env TF_CUDA_PATHS="${cuda_build_deps_joined},${cudnn},${nccl}" build --action_env TF_CUDA_VERSION="${lib.versions.majorMinor cudaVersion}" build --action_env TF_CUDNN_VERSION="${lib.versions.major cudnn.version}" build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="${builtins.concatStringsSep "," cudaFlags.realArches}" '' + # Note that upstream conditions this on `wheel_cpu == "x86_64"`. We just # rely on `effectiveStdenv.hostPlatform.avxSupport` instead. So far so # good. See https://github.com/google/jax/blob/b9824d7de3cb30f1df738cc42e486db3e9d915ff/build/build.py#L322 # for upstream's version. lib.optionalString (effectiveStdenv.hostPlatform.avxSupport && effectiveStdenv.hostPlatform.isUnix) '' build --config=avx_posix '' + lib.optionalString mklSupport '' build --config=mkl_open_source_only '' + '' CFG ''; </ins> # Make sure Bazel knows about our configuration flags during fetching so that the # relevant dependencies can be downloaded. bazelFlags = [ "-c opt" <del> ] ++ lib.optionals stdenv.cc.isClang [ </del> <ins> # See https://bazel.build/external/advanced#overriding-repositories for # information on --override_repository flag. "--override_repository=xla=${xla}" ] ++ lib.optionals effectiveStdenv.cc.isClang [ </ins> # bazel depends on the compiler frontend automatically selecting these flags based on file # extension but our clang doesn't. # https://github.com/NixOS/nixpkgs/issues/150655 <del> "--cxxopt=-x" "--cxxopt=c++" "--host_cxxopt=-x" "--host_cxxopt=c++" </del> <ins> "--cxxopt=-x" "--cxxopt=c++" "--host_cxxopt=-x" "--host_cxxopt=c++" </ins> ]; # We intentionally overfetch so we can share the fetch derivation across all the different configurations