2024-06-05 15:53:02 +00:00
{
stdenv ,
bazel_5 ,
buildBazelPackage ,
lib ,
fetchFromGitHub ,
symlinkJoin ,
2024-07-31 10:19:44 +00:00
addDriverRunpath ,
2024-06-05 15:53:02 +00:00
fetchpatch ,
fetchzip ,
linkFarm ,
# Python deps
buildPythonPackage ,
pythonAtLeast ,
pythonOlder ,
python ,
# Python libraries
numpy ,
tensorboard ,
abseil-cpp ,
absl-py ,
packaging ,
setuptools ,
wheel ,
keras-preprocessing ,
google-pasta ,
opt-einsum ,
astunparse ,
h5py ,
termcolor ,
grpcio ,
six ,
wrapt ,
protobuf-python ,
tensorflow-estimator-bin ,
dill ,
flatbuffers-python ,
portpicker ,
tblib ,
typing-extensions ,
# Common deps
git ,
pybind11 ,
which ,
binutils ,
glibcLocales ,
cython ,
perl ,
# Common libraries
jemalloc ,
mpi ,
gast ,
grpc ,
sqlite ,
boringssl ,
jsoncpp ,
nsync ,
curl ,
snappy ,
flatbuffers-core ,
icu ,
double-conversion ,
libpng ,
libjpeg_turbo ,
giflib ,
protobuf-core ,
# Upstream by default includes cuda support since tensorflow 1.15. We could do
# that in nix as well. It would make some things easier and less confusing, but
# it would also make the default tensorflow package unfree. See
# https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
config ,
cudaSupport ? config . cudaSupport ,
cudaPackages ,
cudaCapabilities ? cudaPackages . cudaFlags . cudaCapabilities ,
mklSupport ? false ,
mkl ,
tensorboardSupport ? true ,
# XLA without CUDA is broken
xlaSupport ? cudaSupport ,
sse42Support ? stdenv . hostPlatform . sse4_2Support ,
avx2Support ? stdenv . hostPlatform . avx2Support ,
fmaSupport ? stdenv . hostPlatform . fmaSupport ,
# Darwin deps
Foundation ,
Security ,
cctools ,
llvmPackages ,
2020-04-24 23:36:52 +00:00
} :
2022-04-15 01:41:22 +00:00
let
2023-03-08 16:32:21 +00:00
originalStdenv = stdenv ;
in
let
# Tensorflow looks at many toolchain-related variables which may diverge.
#
# Toolchain for cuda-enabled builds.
# We want to achieve two things:
# 1. NVCC should use a compatible back-end (e.g. gcc11 for cuda11)
# 2. Normal C++ files should be compiled with the same toolchain,
# to avoid potential weird dynamic linkage errors at runtime.
# This may not be necessary though
#
# Toolchain for Darwin:
# clang 7 fails to emit a symbol for
# __ZN4llvm11SmallPtrSetIPKNS_10AllocaInstELj8EED1Ev in any of the
# translation units, so the build fails at link time
stdenv =
2024-06-05 15:53:02 +00:00
if cudaSupport then
cudaPackages . backendStdenv
else if originalStdenv . isDarwin then
llvmPackages . stdenv
else
originalStdenv ;
2024-05-15 15:35:15 +00:00
inherit ( cudaPackages ) cudatoolkit nccl ;
2023-11-16 04:20:00 +00:00
# use compatible cuDNN (https://www.tensorflow.org/install/source#gpu)
# cudaPackages.cudnn led to this:
# https://github.com/tensorflow/tensorflow/issues/60398
2024-01-02 11:29:13 +00:00
cudnnAttribute = " c u d n n _ 8 _ 6 " ;
2024-07-27 06:49:29 +00:00
cudnnMerged = symlinkJoin {
name = " c u d n n - m e r g e d " ;
paths = [
( lib . getDev cudaPackages . ${ cudnnAttribute } )
( lib . getLib cudaPackages . ${ cudnnAttribute } )
] ;
} ;
2023-11-16 04:20:00 +00:00
gentoo-patches = fetchzip {
url = " h t t p s : / / d e v . g e n t o o . o r g / ~ p e r f i n i o n / p a t c h e s / t e n s o r f l o w - p a t c h e s - 2 . 1 2 . 0 . t a r . b z 2 " ;
hash = " s h a 2 5 6 - S C R X / 5 / z M L 7 L m K E P J k c M 5 T e b e z 9 v v / g m E 4 x h T / j y q W s = " ;
} ;
protobuf-extra = linkFarm " p r o t o b u f - e x t r a " [
2024-06-05 15:53:02 +00:00
{
name = " i n c l u d e " ;
path = protobuf-core . src ;
}
2023-11-16 04:20:00 +00:00
] ;
2020-04-24 23:36:52 +00:00
2021-01-15 22:18:51 +00:00
withTensorboard = ( pythonOlder " 3 . 6 " ) || tensorboardSupport ;
2020-04-24 23:36:52 +00:00
2024-07-27 06:49:29 +00:00
cudaComponents = with cudaPackages ; [
( cuda_nvcc . __spliced . buildHost or cuda_nvcc )
( cuda_nvprune . __spliced . buildHost or cuda_nvprune )
cuda_cccl # block_load.cuh
cuda_cudart # cuda.h
cuda_cupti # cupti.h
cuda_nvcc # See https://github.com/google/jax/issues/19811
cuda_nvml_dev # nvml.h
cuda_nvtx # nvToolsExt.h
libcublas # cublas_api.h
libcufft # cufft.h
libcurand # curand.h
libcusolver # cusolver_common.h
libcusparse # cusparse.h
] ;
cudatoolkitDevMerged = symlinkJoin {
name = " c u d a - ${ cudaPackages . cudaVersion } - d e v - m e r g e d " ;
paths = lib . concatMap ( p : [
( lib . getBin p )
( lib . getDev p )
( lib . getLib p )
( lib . getOutput " s t a t i c " p ) # Makes for a very fat closure
] ) cudaComponents ;
2020-04-24 23:36:52 +00:00
} ;
2023-03-08 16:32:21 +00:00
# Tensorflow expects bintools at hard-coded paths, e.g. /usr/bin/ar
# The only way to overcome that is to set GCC_HOST_COMPILER_PREFIX,
# but that path must contain cc as well, so we merge them
2020-04-24 23:36:52 +00:00
cudatoolkit_cc_joined = symlinkJoin {
2023-03-08 16:32:21 +00:00
name = " ${ stdenv . cc . name } - m e r g e d " ;
2020-04-24 23:36:52 +00:00
paths = [
2023-03-08 16:32:21 +00:00
stdenv . cc
2020-04-24 23:36:52 +00:00
binutils . bintools # for ar, dwp, nm, objcopy, objdump, strip
] ;
} ;
# Needed for _some_ system libraries, grep INCLUDEDIR.
includes_joined = symlinkJoin {
name = " t e n s o r f l o w - d e p s - m e r g e d " ;
2024-06-05 15:53:02 +00:00
paths = [ jsoncpp ] ;
2020-04-24 23:36:52 +00:00
} ;
tfFeature = x : if x then " 1 " else " 0 " ;
2023-11-16 04:20:00 +00:00
version = " 2 . 1 3 . 0 " ;
2024-01-02 11:29:13 +00:00
format = " s e t u p t o o l s " ;
2023-02-16 17:41:37 +00:00
variant = lib . optionalString cudaSupport " - g p u " ;
2020-04-24 23:36:52 +00:00
pname = " t e n s o r f l o w ${ variant } " ;
2024-06-05 15:53:02 +00:00
pythonEnv = python . withPackages ( _ : [
# python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
# This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
absl-py
astunparse
dill
flatbuffers-python
gast
google-pasta
grpcio
h5py
keras-preprocessing
numpy
opt-einsum
packaging
protobuf-python
setuptools
six
tblib
tensorboard
tensorflow-estimator-bin
termcolor
typing-extensions
wheel
wrapt
2020-04-24 23:36:52 +00:00
] ) ;
2021-12-06 16:07:01 +00:00
rules_cc_darwin_patched = stdenv . mkDerivation {
name = " r u l e s _ c c - ${ pname } - ${ version } " ;
src = _bazel-build . deps ;
prePatch = " p u s h d r u l e s _ c c " ;
patches = [
# https://github.com/bazelbuild/rules_cc/issues/122
( fetchpatch {
name = " t e n s o r f l o w - r u l e s _ c c - l i b t o o l - p a t h . p a t c h " ;
url = " h t t p s : / / g i t h u b . c o m / b a z e l b u i l d / r u l e s _ c c / c o m m i t / 8 c 4 2 7 a b 3 0 b f 2 1 3 6 3 0 d c 3 b c e 9 d 2 e 9 a 0 e 2 9 d 1 7 8 7 d b . d i f f " ;
2023-03-15 16:39:30 +00:00
hash = " s h a 2 5 6 - C 4 v 6 H Y 5 + j m 0 A C U Z 5 8 g B P V e j C Y C Z f u z Y K l H Z 0 m 2 q D H C k = " ;
2021-12-06 16:07:01 +00:00
} )
# https://github.com/bazelbuild/rules_cc/pull/124
( fetchpatch {
name = " t e n s o r f l o w - r u l e s _ c c - i n s t a l l _ n a m e _ t o o l - p a t h . p a t c h " ;
url = " h t t p s : / / g i t h u b . c o m / b a z e l b u i l d / r u l e s _ c c / c o m m i t / 1 5 6 4 9 7 d c 8 9 1 0 0 d b 8 a 3 f 5 7 b 2 3 c 6 3 7 2 4 7 5 9 d 4 3 1 d 0 5 . d i f f " ;
2023-03-15 16:39:30 +00:00
hash = " s h a 2 5 6 - N E S 1 K e Q m M i U J Q V o V 6 d S 4 Y G R x x k Z E j O p F S C y O q 9 H Z Y O 0 = " ;
2021-12-06 16:07:01 +00:00
} )
] ;
postPatch = " p o p d " ;
dontConfigure = true ;
dontBuild = true ;
installPhase = ''
runHook preInstall
mv rules_cc / " $ o u t "
runHook postInstall
'' ;
} ;
llvm-raw_darwin_patched = stdenv . mkDerivation {
name = " l l v m - r a w - ${ pname } - ${ version } " ;
src = _bazel-build . deps ;
prePatch = " p u s h d l l v m - r a w " ;
patches = [
# Fix a vendored config.h that requires the 10.13 SDK
./llvm_bazel_fix_macos_10_12_sdk.patch
] ;
postPatch = ''
touch { BUILD , WORKSPACE }
popd
'' ;
dontConfigure = true ;
dontBuild = true ;
installPhase = ''
runHook preInstall
mv llvm-raw / " $ o u t "
runHook postInstall
'' ;
} ;
2024-06-05 15:53:02 +00:00
bazel-build =
if stdenv . isDarwin then
_bazel-build . overrideAttrs ( prev : {
bazelFlags = prev . bazelFlags ++ [
" - - o v e r r i d e _ r e p o s i t o r y = r u l e s _ c c = ${ rules_cc_darwin_patched } "
" - - o v e r r i d e _ r e p o s i t o r y = l l v m - r a w = ${ llvm-raw_darwin_patched } "
] ;
preBuild = ''
export AR = " ${ cctools } / b i n / l i b t o o l "
'' ;
} )
else
_bazel-build ;
2021-12-06 16:07:01 +00:00
2023-03-08 16:32:21 +00:00
_bazel-build = buildBazelPackage . override { inherit stdenv ; } {
2020-04-24 23:36:52 +00:00
name = " ${ pname } - ${ version } " ;
2022-08-12 12:06:08 +00:00
bazel = bazel_5 ;
2020-04-24 23:36:52 +00:00
src = fetchFromGitHub {
owner = " t e n s o r f l o w " ;
repo = " t e n s o r f l o w " ;
2023-02-02 18:25:31 +00:00
rev = " r e f s / t a g s / v ${ version } " ;
2023-11-16 04:20:00 +00:00
hash = " s h a 2 5 6 - R q 5 p A V m x l W B V n p h 2 0 f k A w b f y + i u B N l f F y 1 4 p o D P d 5 h 0 = " ;
2020-04-24 23:36:52 +00:00
} ;
# On update, it can be useful to steal the changes from gentoo
# https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
nativeBuildInputs = [
2024-06-05 15:53:02 +00:00
which
pythonEnv
cython
perl
protobuf-core
protobuf-extra
2024-07-31 10:19:44 +00:00
] ++ lib . optional cudaSupport addDriverRunpath ;
2020-04-24 23:36:52 +00:00
2024-06-05 15:53:02 +00:00
buildInputs =
[
jemalloc
mpi
glibcLocales
git
# libs taken from system through the TF_SYS_LIBS mechanism
abseil-cpp
boringssl
curl
double-conversion
flatbuffers-core
giflib
grpc
# Necessary to fix the "`GLIBCXX_3.4.30' not found" error
( icu . override { inherit stdenv ; } )
jsoncpp
libjpeg_turbo
libpng
( pybind11 . overridePythonAttrs ( _ : {
inherit stdenv ;
} ) )
snappy
sqlite
]
++ lib . optionals cudaSupport [
cudatoolkit
2024-07-27 06:49:29 +00:00
cudnnMerged
2024-06-05 15:53:02 +00:00
]
++ lib . optionals mklSupport [ mkl ]
++ lib . optionals stdenv . isDarwin [
Foundation
Security
]
++ lib . optionals ( ! stdenv . isDarwin ) [ nsync ] ;
2020-04-24 23:36:52 +00:00
# arbitrarily set to the current latest bazel version, overly careful
TF_IGNORE_MAX_BAZEL_VERSION = true ;
2022-12-02 08:20:57 +00:00
LIBTOOL = lib . optionalString stdenv . isDarwin " ${ cctools } / b i n / l i b t o o l " ;
2020-04-24 23:36:52 +00:00
# Take as many libraries from the system as possible. Keep in sync with
# list of valid syslibs in
# https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
2024-06-05 15:53:02 +00:00
TF_SYSTEM_LIBS = lib . concatStringsSep " , " (
[
" a b s l _ p y "
" a s t o r _ a r c h i v e "
" a s t u n p a r s e _ a r c h i v e "
" b o r i n g s s l "
" c o m _ g o o g l e _ a b s l "
# Not packaged in nixpkgs
# "com_github_googleapis_googleapis"
# "com_github_googlecloudplatform_google_cloud_cpp"
" c o m _ g i t h u b _ g r p c _ g r p c "
" c o m _ g o o g l e _ p r o t o b u f "
# Fails with the error: external/org_tensorflow/tensorflow/core/profiler/utils/tf_op_utils.cc:46:49: error: no matching function for call to 're2::RE2::FullMatch(absl::lts_2020_02_25::string_view&, re2::RE2&)'
# "com_googlesource_code_re2"
" c u r l "
" c y t h o n "
" d i l l _ a r c h i v e "
" d o u b l e _ c o n v e r s i o n "
" f l a t b u f f e r s "
" f u n c t o o l s 3 2 _ a r c h i v e "
" g a s t _ a r c h i v e "
" g i f "
" h w l o c "
" i c u "
" j s o n c p p _ g i t "
" l i b j p e g _ t u r b o "
" n a s m "
" o p t _ e i n s u m _ a r c h i v e "
" o r g _ s q l i t e "
" p a s t a "
" p n g "
" p y b i n d 1 1 "
" s i x _ a r c h i v e "
" s n a p p y "
" t b l i b _ a r c h i v e "
" t e r m c o l o r _ a r c h i v e "
" t y p i n g _ e x t e n s i o n s _ a r c h i v e "
" w r a p t "
" z l i b "
]
++ lib . optionals ( ! stdenv . isDarwin ) [
" n s y n c " # fails to build on darwin
]
) ;
2020-04-24 23:36:52 +00:00
INCLUDEDIR = " ${ includes_joined } / i n c l u d e " ;
2021-12-19 01:06:50 +00:00
# This is needed for the Nix-provided protobuf dependency to work,
# as otherwise the rule `link_proto_files` tries to create the links
# to `/usr/include/...` which results in build failures.
PROTOBUF_INCLUDE_PATH = " ${ protobuf-core } / i n c l u d e " ;
2020-04-24 23:36:52 +00:00
PYTHON_BIN_PATH = pythonEnv . interpreter ;
TF_NEED_GCP = true ;
TF_NEED_HDFS = true ;
TF_ENABLE_XLA = tfFeature xlaSupport ;
CC_OPT_FLAGS = " " ;
# https://github.com/tensorflow/tensorflow/issues/14454
TF_NEED_MPI = tfFeature cudaSupport ;
TF_NEED_CUDA = tfFeature cudaSupport ;
2024-07-27 06:49:29 +00:00
TF_CUDA_PATHS = lib . optionalString cudaSupport " ${ cudatoolkitDevMerged } , ${ cudnnMerged } , ${ lib . getLib nccl } " ;
2023-03-08 16:32:21 +00:00
TF_CUDA_COMPUTE_CAPABILITIES = lib . concatStringsSep " , " cudaCapabilities ;
# Needed even when we override stdenv: e.g. for ar
2020-04-24 23:36:52 +00:00
GCC_HOST_COMPILER_PREFIX = lib . optionalString cudaSupport " ${ cudatoolkit_cc_joined } / b i n " ;
2023-03-08 16:32:21 +00:00
GCC_HOST_COMPILER_PATH = lib . optionalString cudaSupport " ${ cudatoolkit_cc_joined } / b i n / c c " ;
2020-04-24 23:36:52 +00:00
2023-11-16 04:20:00 +00:00
patches = [
" ${ gentoo-patches } / 0 0 0 2 - s y s t e m l i b - L a t e s t - a b s l - L T S - h a s - s p l i t - c o r d - l i b s . p a t c h "
" ${ gentoo-patches } / 0 0 0 5 - s y s t e m l i b - U p d a t e s - f o r - A b s e i l - 2 0 2 2 0 6 2 3 - L T S . p a t c h "
" ${ gentoo-patches } / 0 0 0 7 - s y s t e m l i b s - A d d - w e l l _ k n o w n _ t y p e s _ p y _ p b 2 - t a r g e t . p a t c h "
# https://github.com/conda-forge/tensorflow-feedstock/pull/329/commits/0a63c5a962451b4da99a9948323d8b3ed462f461
( fetchpatch {
name = " f i x - l a y o u t - p r o t o - d u p l i c a t e - l o a d i n g . p a t c h " ;
url = " h t t p s : / / r a w . g i t h u b u s e r c o n t e n t . c o m / c o n d a - f o r g e / t e n s o r f l o w - f e e d s t o c k / 0 a 6 3 c 5 a 9 6 2 4 5 1 b 4 d a 9 9 a 9 9 4 8 3 2 3 d 8 b 3 e d 4 6 2 f 4 6 1 / r e c i p e / p a t c h e s / 0 0 0 1 - O m i t - l i n k i n g - t o - l a y o u t _ p r o t o _ c c - i f - p r o t o b u f - l i n k a g e - . p a t c h " ;
hash = " s h a 2 5 6 - / 7 b u V 6 D i n K n r g f q b e 7 K K S h 9 r C e b e Q d X v 2 U j + X g / 0 8 3 w = " ;
} )
./com_google_absl_add_log.patch
./absl_py_argparse_flags.patch
./protobuf_python.patch
./pybind11_protobuf_python_runtime_dep.patch
./pybind11_protobuf_newer_version.patch
2024-06-05 15:53:02 +00:00
] ++ lib . optionals ( stdenv . hostPlatform . system == " a a r c h 6 4 - d a r w i n " ) [ ./absl_to_std.patch ] ;
2023-11-16 04:20:00 +00:00
2024-06-05 15:53:02 +00:00
postPatch =
''
# bazel 3.3 should work just as well as bazel 3.1
rm - f . bazelversion
patchShebangs .
''
+ lib . optionalString ( stdenv . hostPlatform . system == " x 8 6 _ 6 4 - d a r w i n " ) ''
cat $ { ./com_google_absl_fix_macos.patch } > > third_party/absl/com_google_absl_fix_mac_and_nvcc_build.patch
''
+ lib . optionalString ( ! withTensorboard ) ''
# Tensorboard pulls in a bunch of dependencies, some of which may
# include security vulnerabilities. So we make it optional.
# https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
sed - i ' /tensorboard ~ = /d ' tensorflow/tools/pip_package/setup.py
'' ;
2020-04-24 23:36:52 +00:00
2021-01-09 10:05:03 +00:00
# https://github.com/tensorflow/tensorflow/pull/39470
2023-03-04 12:14:45 +00:00
env . NIX_CFLAGS_COMPILE = toString [ " - W n o - s t r i n g o p - t r u n c a t i o n " ] ;
2021-01-09 10:05:03 +00:00
2024-06-05 15:53:02 +00:00
preConfigure =
let
opt_flags =
[ ]
++ lib . optionals sse42Support [ " - m s s e 4 . 2 " ]
++ lib . optionals avx2Support [ " - m a v x 2 " ]
++ lib . optionals fmaSupport [ " - m f m a " ] ;
in
''
patchShebangs configure
# dummy ldconfig
mkdir dummy-ldconfig
echo " # ! ${ stdenv . shell } " > dummy-ldconfig/ldconfig
chmod + x dummy-ldconfig/ldconfig
export PATH = " $ P W D / d u m m y - l d c o n f i g : $ P A T H "
export PYTHON_LIB_PATH = " $ N I X _ B U I L D _ T O P / s i t e - p a c k a g e s "
export CC_OPT_FLAGS = " ${ lib . concatStringsSep " " opt_flags } "
mkdir - p " $ P Y T H O N _ L I B _ P A T H "
# To avoid mixing Python 2 and Python 3
unset PYTHONPATH
'' ;
2020-04-24 23:36:52 +00:00
configurePhase = ''
runHook preConfigure
./configure
runHook postConfigure
'' ;
hardeningDisable = [ " f o r m a t " ] ;
2024-06-05 15:53:02 +00:00
bazelBuildFlags =
[
" - - c o n f i g = o p t " # optimize using the flags set in the configure phase
]
++ lib . optionals stdenv . cc . isClang [
" - - c x x o p t = - x "
" - - c x x o p t = c + + "
" - - h o s t _ c x x o p t = - x "
" - - h o s t _ c x x o p t = c + + "
# workaround for https://github.com/bazelbuild/bazel/issues/15359
" - - s p a w n _ s t r a t e g y = s a n d b o x e d "
]
++ lib . optionals ( mklSupport ) [ " - - c o n f i g = m k l " ] ;
bazelTargets = [
" / / t e n s o r f l o w / t o o l s / p i p _ p a c k a g e : b u i l d _ p i p _ p a c k a g e / / t e n s o r f l o w / t o o l s / l i b _ p a c k a g e : l i b t e n s o r f l o w "
] ;
2020-04-24 23:36:52 +00:00
2020-10-07 09:15:18 +00:00
removeRulesCC = false ;
2021-01-15 22:18:51 +00:00
# Without this Bazel complaints about sandbox violations.
dontAddBazelOpts = true ;
2020-10-07 09:15:18 +00:00
2020-04-24 23:36:52 +00:00
fetchAttrs = {
2024-06-05 15:53:02 +00:00
sha256 =
{
x86_64-linux =
if cudaSupport then
" s h a 2 5 6 - 5 V F M N H e L r U x W 5 R T r 6 E h T 3 p a y 9 n W J 5 J k Z T G i r D d s 5 Q k U = "
else
" s h a 2 5 6 - K z g W V 6 9 B t r 8 4 F d w Q 5 J I 2 n Q E s q i P g 1 / + T W d b w 5 b m x X O E = " ;
aarch64-linux =
if cudaSupport then
" s h a 2 5 6 - t y 5 + 5 1 B w H W E 1 x R 4 / 0 W c W T p 6 0 8 N z S A S / i i y N + 9 z x 7 / w I = "
else
" s h a 2 5 6 - 9 b t X r N H q d 7 2 0 o X T P D h S m F i d v 5 i a Z R L j C V X 8 o p m r M j X k = " ;
x86_64-darwin = " s h a 2 5 6 - g q b 0 3 k B 0 z 2 p Z Q 6 m 1 f y R p 1 / N b t 8 A V V H W p O J S e Z N C L c 4 w = " ;
aarch64-darwin = " s h a 2 5 6 - W d g A a F Z U + e P w W k V B h L z j l N T 7 E L f G H O T a M d a f c A M D 5 y o = " ;
}
. ${ stdenv . hostPlatform . system } or ( throw " u n s u p p o r t e d s y s t e m ${ stdenv . hostPlatform . system } " ) ;
2020-04-24 23:36:52 +00:00
} ;
buildAttrs = {
2024-06-05 15:53:02 +00:00
outputs = [
" o u t "
" p y t h o n "
] ;
2020-04-24 23:36:52 +00:00
2023-11-16 04:20:00 +00:00
# need to rebuild schemas since we use a different flatbuffers version
2020-04-24 23:36:52 +00:00
preBuild = ''
2023-11-16 04:20:00 +00:00
( cd tensorflow/lite/schema ; $ { flatbuffers-core } /bin/flatc - - gen-object-api - c schema . fbs )
( cd tensorflow/lite/schema ; $ { flatbuffers-core } /bin/flatc - - gen-object-api - c conversion_metadata . fbs )
( cd tensorflow/lite/acceleration/configuration ; $ { flatbuffers-core } /bin/flatc - o configuration . fbs - - proto configuration . proto )
sed - i s , tflite . proto , tflite , g tensorflow/lite/acceleration/configuration/configuration.fbs/configuration.fbs
( cd tensorflow/lite/acceleration/configuration ; $ { flatbuffers-core } /bin/flatc - - gen-compare - - gen-object-api - c configuration.fbs/configuration.fbs )
cp - r tensorflow/lite/acceleration/configuration/configuration.fbs tensorflow/lite/experimental/acceleration/configuration
( cd tensorflow/lite/experimental/acceleration/configuration ; $ { flatbuffers-core } /bin/flatc - c configuration.fbs/configuration.fbs )
( cd tensorflow/lite/delegates/gpu/cl ; $ { flatbuffers-core } /bin/flatc - c compiled_program_cache . fbs )
( cd tensorflow/lite/delegates/gpu/cl ; $ { flatbuffers-core } /bin/flatc - I $ NIX_BUILD_TOP/source - c serialization . fbs )
( cd tensorflow/lite/delegates/gpu/common ; $ { flatbuffers-core } /bin/flatc - I $ NIX_BUILD_TOP/source - c gpu_model . fbs )
( cd tensorflow/lite/delegates/gpu/common/task ; $ { flatbuffers-core } /bin/flatc - c serialization_base . fbs )
2020-04-24 23:36:52 +00:00
patchShebangs .
'' ;
installPhase = ''
mkdir - p " $ o u t "
tar - xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz - C " $ o u t "
# Write pkgconfig file.
mkdir " $ o u t / l i b / p k g c o n f i g "
cat > " $ o u t / l i b / p k g c o n f i g / t e n s o r f l o w . p c " < < EOF
Name : TensorFlow
Version : $ { version }
Description : Library for computation using data flow graphs for scalable machine learning
Requires :
Libs : - L $ out/lib - ltensorflow
Cflags : - I $ out/include/tensorflow
EOF
# build the source code, then copy it to $python (build_pip_package
# actually builds a symlink farm so we must dereference them).
bazel-bin/tensorflow/tools/pip_package/build_pip_package - - src " $ P W D / d i s t "
cp - Lr " $ P W D / d i s t " " $ p y t h o n "
'' ;
postFixup = lib . optionalString cudaSupport ''
find $ out - type f \ ( - name ' * . so' - or - name ' * . so . * ' \ ) | while read lib ; do
2024-07-31 10:19:44 +00:00
addDriverRunpath " $ l i b "
2020-04-24 23:36:52 +00:00
done
'' ;
2021-04-26 19:14:03 +00:00
2024-06-05 15:53:02 +00:00
requiredSystemFeatures = [ " b i g - p a r a l l e l " ] ;
2020-04-24 23:36:52 +00:00
} ;
2024-06-05 15:53:02 +00:00
meta =
{
badPlatforms = lib . optionals cudaSupport lib . platforms . darwin ;
changelog = " h t t p s : / / g i t h u b . c o m / t e n s o r f l o w / t e n s o r f l o w / r e l e a s e s / t a g / v ${ version } " ;
description = " C o m p u t a t i o n u s i n g d a t a f l o w g r a p h s f o r s c a l a b l e m a c h i n e l e a r n i n g " ;
homepage = " h t t p : / / t e n s o r f l o w . o r g " ;
2024-09-19 14:19:46 +00:00
license = lib . licenses . asl20 ;
maintainers = with lib . maintainers ; [ abbradar ] ;
platforms = with lib . platforms ; linux ++ darwin ;
2024-06-05 15:53:02 +00:00
broken =
stdenv . isDarwin
|| ! ( xlaSupport -> cudaSupport )
|| ! ( cudaSupport -> builtins . hasAttr cudnnAttribute cudaPackages )
|| ! ( cudaSupport -> cudaPackages ? cudatoolkit ) ;
}
// lib . optionalAttrs stdenv . isDarwin {
timeout = 86400 ; # 24 hours
maxSilent = 14400 ; # 4h, double the default of 7200s
} ;
2020-04-24 23:36:52 +00:00
} ;
2024-06-05 15:53:02 +00:00
in
buildPythonPackage {
2020-04-24 23:36:52 +00:00
inherit version pname ;
2024-05-15 15:35:15 +00:00
disabled = pythonAtLeast " 3 . 1 2 " ;
2020-04-24 23:36:52 +00:00
src = bazel-build . python ;
2021-12-06 16:07:01 +00:00
# Adjust dependency requirements:
2023-02-02 18:25:31 +00:00
# - Drop tensorflow-io dependency until we get it to build
2022-06-26 10:26:21 +00:00
# - Relax flatbuffers and gast version requirements
2021-12-06 16:07:01 +00:00
# - The purpose of python3Packages.libclang is not clear at the moment and we don't have it packaged yet
# - keras and tensorlow-io-gcs-filesystem will be considered as optional for now.
postPatch = ''
sed - i setup . py \
2023-02-02 18:25:31 +00:00
- e ' /tensorflow-io-gcs-filesystem / , + 1 d' \
2022-06-26 10:26:21 +00:00
- e " s / ' f l a t b u f f e r s [ ^ ' ] * ' , / ' f l a t b u f f e r s ' , / " \
2021-12-06 16:07:01 +00:00
- e " s / ' g a s t [ ^ ' ] * ' , / ' g a s t ' , / " \
- e " / ' l i b c l a n g [ ^ ' ] * ' , / d " \
2022-06-26 10:26:21 +00:00
- e " / ' k e r a s [ ^ ' ] * ' ) \? , / d " \
2022-09-30 11:47:45 +00:00
- e " / ' t e n s o r f l o w - i o - g c s - f i l e s y s t e m [ ^ ' ] * ' , / d " \
- e " s / ' p r o t o b u f [ ^ ' ] * ' , / ' p r o t o b u f ' , / " \
2021-12-06 16:07:01 +00:00
'' ;
2020-04-24 23:36:52 +00:00
# Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
2022-04-15 01:41:22 +00:00
# and the propagated input tensorboard, which causes environment collisions.
2020-04-24 23:36:52 +00:00
# Another possibility would be to have tensorboard only in the buildInputs
# https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
postInstall = ''
rm $ out/bin/tensorboard
'' ;
setupPyGlobalFlags = [ " - - p r o j e c t _ n a m e ${ pname } " ] ;
# tensorflow/tools/pip_package/setup.py
propagatedBuildInputs = [
absl-py
2023-11-16 04:20:00 +00:00
abseil-cpp
2021-01-15 22:18:51 +00:00
astunparse
flatbuffers-python
2020-04-24 23:36:52 +00:00
gast
google-pasta
2021-01-15 22:18:51 +00:00
grpcio
h5py
2020-04-24 23:36:52 +00:00
keras-preprocessing
numpy
2021-01-15 22:18:51 +00:00
opt-einsum
2022-06-26 10:26:21 +00:00
packaging
2021-12-19 01:06:50 +00:00
protobuf-python
2021-01-15 22:18:51 +00:00
six
2023-02-02 18:25:31 +00:00
tensorflow-estimator-bin
2020-04-24 23:36:52 +00:00
termcolor
2021-01-15 22:18:51 +00:00
typing-extensions
2020-04-24 23:36:52 +00:00
wrapt
2024-06-05 15:53:02 +00:00
] ++ lib . optionals withTensorboard [ tensorboard ] ;
2020-04-24 23:36:52 +00:00
2024-07-31 10:19:44 +00:00
nativeBuildInputs = lib . optionals cudaSupport [ addDriverRunpath ] ;
2020-04-24 23:36:52 +00:00
postFixup = lib . optionalString cudaSupport ''
find $ out - type f \ ( - name ' * . so' - or - name ' * . so . * ' \ ) | while read lib ; do
2024-07-31 10:19:44 +00:00
addDriverRunpath " $ l i b "
2020-10-07 09:15:18 +00:00
2024-07-27 06:49:29 +00:00
patchelf - - set-rpath " ${ cudatoolkit } / l i b : ${ cudatoolkit . lib } / l i b : ${ cudnnMerged } / l i b : ${ lib . getLib nccl } / l i b : $ ( p a t c h e l f - - p r i n t - r p a t h " $ lib " ) " " $ l i b "
2020-04-24 23:36:52 +00:00
done
'' ;
# Actual tests are slow and impure.
# TODO try to run them anyway
# TODO better test (files in tensorflow/tools/ci_build/builds/*test)
2022-04-15 01:41:22 +00:00
# TEST_PACKAGES in tensorflow/tools/pip_package/setup.py
2023-02-02 18:25:31 +00:00
nativeCheckInputs = [
2022-04-15 01:41:22 +00:00
dill
portpicker
tblib
] ;
2020-04-24 23:36:52 +00:00
checkPhase = ''
$ { python . interpreter } < < EOF
# A simple "Hello world"
import tensorflow as tf
hello = tf . constant ( " H e l l o , w o r l d ! " )
tf . print ( hello )
tf . random . set_seed ( 0 )
2023-11-16 04:20:00 +00:00
width = 512
choice = 48
t_in = tf . Variable ( tf . random . uniform ( shape = [ width ] ) )
with tf . GradientTape ( ) as tape :
t_out = tf . slice ( tf . nn . softmax ( t_in ) , [ choice ] , [ 1 ] )
diff = tape . gradient ( t_out , t_in )
assert ( 0 < tf . reduce_min ( tf . slice ( diff , [ choice ] , [ 1 ] ) ) )
assert ( 0 > tf . reduce_max ( tf . slice ( diff , [ 1 ] , [ choice - 1 ] ) ) )
2020-04-24 23:36:52 +00:00
EOF
'' ;
# Regression test for #77626 removed because not more `tensorflow.contrib`.
2020-05-29 06:06:01 +00:00
passthru = {
deps = bazel-build . deps ;
libtensorflow = bazel-build . out ;
} ;
2020-04-24 23:36:52 +00:00
inherit ( bazel-build ) meta ;
}
