tjh.dev
1{
2 stdenv,
3 bazel_5,
4 buildBazelPackage,
5 lib,
6 fetchFromGitHub,
7 symlinkJoin,
8 addDriverRunpath,
9 fetchpatch,
10 fetchzip,
11 linkFarm,
12 # Python deps
13 buildPythonPackage,
14 pythonAtLeast,
15 pythonOlder,
16 python,
17 # Python libraries
18 numpy,
19 tensorboard,
20 abseil-cpp,
21 absl-py,
22 packaging,
23 setuptools,
24 wheel,
25 keras-preprocessing,
26 google-pasta,
27 opt-einsum,
28 astunparse,
29 h5py,
30 termcolor,
31 grpcio,
32 six,
33 wrapt,
34 protobuf-python,
35 tensorflow-estimator-bin,
36 dill,
37 flatbuffers-python,
38 portpicker,
39 tblib,
40 typing-extensions,
41 # Common deps
42 git,
43 pybind11,
44 which,
45 binutils,
46 glibcLocales,
47 cython,
48 perl,
49 # Common libraries
50 jemalloc,
51 mpi,
52 gast,
53 grpc,
54 sqlite,
55 boringssl,
56 jsoncpp,
57 nsync,
58 curl,
59 snappy,
60 flatbuffers-core,
61 icu,
62 double-conversion,
63 libpng,
64 libjpeg_turbo,
65 giflib,
66 protobuf-core,
67 # Upstream by default includes cuda support since tensorflow 1.15. We could do
68 # that in nix as well. It would make some things easier and less confusing, but
69 # it would also make the default tensorflow package unfree. See
70 # https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
71 config,
72 cudaSupport ? config.cudaSupport,
73 cudaPackages,
74 cudaCapabilities ? cudaPackages.cudaFlags.cudaCapabilities,
75 mklSupport ? false,
76 mkl,
77 tensorboardSupport ? true,
78 # XLA without CUDA is broken
79 xlaSupport ? cudaSupport,
80 sse42Support ? stdenv.hostPlatform.sse4_2Support,
81 avx2Support ? stdenv.hostPlatform.avx2Support,
82 fmaSupport ? stdenv.hostPlatform.fmaSupport,
83 # Darwin deps
84 Foundation,
85 Security,
86 cctools,
87 llvmPackages,
88}:
89
90let
91 originalStdenv = stdenv;
92in
93let
94 # Tensorflow looks at many toolchain-related variables which may diverge.
95 #
96 # Toolchain for cuda-enabled builds.
97 # We want to achieve two things:
98 # 1. NVCC should use a compatible back-end (e.g. gcc11 for cuda11)
99 # 2. Normal C++ files should be compiled with the same toolchain,
100 # to avoid potential weird dynamic linkage errors at runtime.
101 # This may not be necessary though
102 #
103 # Toolchain for Darwin:
104 # clang 7 fails to emit a symbol for
105 # __ZN4llvm11SmallPtrSetIPKNS_10AllocaInstELj8EED1Ev in any of the
106 # translation units, so the build fails at link time
107 stdenv =
108 if cudaSupport then
109 cudaPackages.backendStdenv
110 else if originalStdenv.isDarwin then
111 llvmPackages.stdenv
112 else
113 originalStdenv;
114 inherit (cudaPackages) cudatoolkit nccl;
115 # use compatible cuDNN (https://www.tensorflow.org/install/source#gpu)
116 # cudaPackages.cudnn led to this:
117 # https://github.com/tensorflow/tensorflow/issues/60398
118 cudnnAttribute = "cudnn_8_6";
119 cudnnMerged = symlinkJoin {
120 name = "cudnn-merged";
121 paths = [
122 (lib.getDev cudaPackages.${cudnnAttribute})
123 (lib.getLib cudaPackages.${cudnnAttribute})
124 ];
125 };
126 gentoo-patches = fetchzip {
127 url = "https://dev.gentoo.org/~perfinion/patches/tensorflow-patches-2.12.0.tar.bz2";
128 hash = "sha256-SCRX/5/zML7LmKEPJkcM5Tebez9vv/gmE4xhT/jyqWs=";
129 };
130 protobuf-extra = linkFarm "protobuf-extra" [
131 {
132 name = "include";
133 path = protobuf-core.src;
134 }
135 ];
136
137 withTensorboard = (pythonOlder "3.6") || tensorboardSupport;
138
139 cudaComponents = with cudaPackages; [
140 (cuda_nvcc.__spliced.buildHost or cuda_nvcc)
141 (cuda_nvprune.__spliced.buildHost or cuda_nvprune)
142 cuda_cccl # block_load.cuh
143 cuda_cudart # cuda.h
144 cuda_cupti # cupti.h
145 cuda_nvcc # See https://github.com/google/jax/issues/19811
146 cuda_nvml_dev # nvml.h
147 cuda_nvtx # nvToolsExt.h
148 libcublas # cublas_api.h
149 libcufft # cufft.h
150 libcurand # curand.h
151 libcusolver # cusolver_common.h
152 libcusparse # cusparse.h
153 ];
154
155 cudatoolkitDevMerged = symlinkJoin {
156 name = "cuda-${cudaPackages.cudaVersion}-dev-merged";
157 paths = lib.concatMap (p: [
158 (lib.getBin p)
159 (lib.getDev p)
160 (lib.getLib p)
161 (lib.getOutput "static" p) # Makes for a very fat closure
162 ]) cudaComponents;
163 };
164
165 # Tensorflow expects bintools at hard-coded paths, e.g. /usr/bin/ar
166 # The only way to overcome that is to set GCC_HOST_COMPILER_PREFIX,
167 # but that path must contain cc as well, so we merge them
168 cudatoolkit_cc_joined = symlinkJoin {
169 name = "${stdenv.cc.name}-merged";
170 paths = [
171 stdenv.cc
172 binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
173 ];
174 };
175
176 # Needed for _some_ system libraries, grep INCLUDEDIR.
177 includes_joined = symlinkJoin {
178 name = "tensorflow-deps-merged";
179 paths = [ jsoncpp ];
180 };
181
182 tfFeature = x: if x then "1" else "0";
183
184 version = "2.13.0";
185 format = "setuptools";
186 variant = lib.optionalString cudaSupport "-gpu";
187 pname = "tensorflow${variant}";
188
189 pythonEnv = python.withPackages (_: [
190 # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
191 # This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
192 absl-py
193 astunparse
194 dill
195 flatbuffers-python
196 gast
197 google-pasta
198 grpcio
199 h5py
200 keras-preprocessing
201 numpy
202 opt-einsum
203 packaging
204 protobuf-python
205 setuptools
206 six
207 tblib
208 tensorboard
209 tensorflow-estimator-bin
210 termcolor
211 typing-extensions
212 wheel
213 wrapt
214 ]);
215
216 rules_cc_darwin_patched = stdenv.mkDerivation {
217 name = "rules_cc-${pname}-${version}";
218
219 src = _bazel-build.deps;
220
221 prePatch = "pushd rules_cc";
222 patches = [
223 # https://github.com/bazelbuild/rules_cc/issues/122
224 (fetchpatch {
225 name = "tensorflow-rules_cc-libtool-path.patch";
226 url = "https://github.com/bazelbuild/rules_cc/commit/8c427ab30bf213630dc3bce9d2e9a0e29d1787db.diff";
227 hash = "sha256-C4v6HY5+jm0ACUZ58gBPVejCYCZfuzYKlHZ0m2qDHCk=";
228 })
229
230 # https://github.com/bazelbuild/rules_cc/pull/124
231 (fetchpatch {
232 name = "tensorflow-rules_cc-install_name_tool-path.patch";
233 url = "https://github.com/bazelbuild/rules_cc/commit/156497dc89100db8a3f57b23c63724759d431d05.diff";
234 hash = "sha256-NES1KeQmMiUJQVoV6dS4YGRxxkZEjOpFSCyOq9HZYO0=";
235 })
236 ];
237 postPatch = "popd";
238
239 dontConfigure = true;
240 dontBuild = true;
241
242 installPhase = ''
243 runHook preInstall
244
245 mv rules_cc/ "$out"
246
247 runHook postInstall
248 '';
249 };
250 llvm-raw_darwin_patched = stdenv.mkDerivation {
251 name = "llvm-raw-${pname}-${version}";
252
253 src = _bazel-build.deps;
254
255 prePatch = "pushd llvm-raw";
256 patches = [
257 # Fix a vendored config.h that requires the 10.13 SDK
258 ./llvm_bazel_fix_macos_10_12_sdk.patch
259 ];
260 postPatch = ''
261 touch {BUILD,WORKSPACE}
262 popd
263 '';
264
265 dontConfigure = true;
266 dontBuild = true;
267
268 installPhase = ''
269 runHook preInstall
270
271 mv llvm-raw/ "$out"
272
273 runHook postInstall
274 '';
275 };
276 bazel-build =
277 if stdenv.isDarwin then
278 _bazel-build.overrideAttrs (prev: {
279 bazelFlags = prev.bazelFlags ++ [
280 "--override_repository=rules_cc=${rules_cc_darwin_patched}"
281 "--override_repository=llvm-raw=${llvm-raw_darwin_patched}"
282 ];
283 preBuild = ''
284 export AR="${cctools}/bin/libtool"
285 '';
286 })
287 else
288 _bazel-build;
289
290 _bazel-build = buildBazelPackage.override { inherit stdenv; } {
291 name = "${pname}-${version}";
292 bazel = bazel_5;
293
294 src = fetchFromGitHub {
295 owner = "tensorflow";
296 repo = "tensorflow";
297 rev = "refs/tags/v${version}";
298 hash = "sha256-Rq5pAVmxlWBVnph20fkAwbfy+iuBNlfFy14poDPd5h0=";
299 };
300
301 # On update, it can be useful to steal the changes from gentoo
302 # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
303
304 nativeBuildInputs = [
305 which
306 pythonEnv
307 cython
308 perl
309 protobuf-core
310 protobuf-extra
311 ] ++ lib.optional cudaSupport addDriverRunpath;
312
313 buildInputs =
314 [
315 jemalloc
316 mpi
317 glibcLocales
318 git
319
320 # libs taken from system through the TF_SYS_LIBS mechanism
321 abseil-cpp
322 boringssl
323 curl
324 double-conversion
325 flatbuffers-core
326 giflib
327 grpc
328 # Necessary to fix the "`GLIBCXX_3.4.30' not found" error
329 (icu.override { inherit stdenv; })
330 jsoncpp
331 libjpeg_turbo
332 libpng
333 (pybind11.overridePythonAttrs (_: {
334 inherit stdenv;
335 }))
336 snappy
337 sqlite
338 ]
339 ++ lib.optionals cudaSupport [
340 cudatoolkit
341 cudnnMerged
342 ]
343 ++ lib.optionals mklSupport [ mkl ]
344 ++ lib.optionals stdenv.isDarwin [
345 Foundation
346 Security
347 ]
348 ++ lib.optionals (!stdenv.isDarwin) [ nsync ];
349
350 # arbitrarily set to the current latest bazel version, overly careful
351 TF_IGNORE_MAX_BAZEL_VERSION = true;
352
353 LIBTOOL = lib.optionalString stdenv.isDarwin "${cctools}/bin/libtool";
354
355 # Take as many libraries from the system as possible. Keep in sync with
356 # list of valid syslibs in
357 # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
358 TF_SYSTEM_LIBS = lib.concatStringsSep "," (
359 [
360 "absl_py"
361 "astor_archive"
362 "astunparse_archive"
363 "boringssl"
364 "com_google_absl"
365 # Not packaged in nixpkgs
366 # "com_github_googleapis_googleapis"
367 # "com_github_googlecloudplatform_google_cloud_cpp"
368 "com_github_grpc_grpc"
369 "com_google_protobuf"
370 # 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&)'
371 # "com_googlesource_code_re2"
372 "curl"
373 "cython"
374 "dill_archive"
375 "double_conversion"
376 "flatbuffers"
377 "functools32_archive"
378 "gast_archive"
379 "gif"
380 "hwloc"
381 "icu"
382 "jsoncpp_git"
383 "libjpeg_turbo"
384 "nasm"
385 "opt_einsum_archive"
386 "org_sqlite"
387 "pasta"
388 "png"
389 "pybind11"
390 "six_archive"
391 "snappy"
392 "tblib_archive"
393 "termcolor_archive"
394 "typing_extensions_archive"
395 "wrapt"
396 "zlib"
397 ]
398 ++ lib.optionals (!stdenv.isDarwin) [
399 "nsync" # fails to build on darwin
400 ]
401 );
402
403 INCLUDEDIR = "${includes_joined}/include";
404
405 # This is needed for the Nix-provided protobuf dependency to work,
406 # as otherwise the rule `link_proto_files` tries to create the links
407 # to `/usr/include/...` which results in build failures.
408 PROTOBUF_INCLUDE_PATH = "${protobuf-core}/include";
409
410 PYTHON_BIN_PATH = pythonEnv.interpreter;
411
412 TF_NEED_GCP = true;
413 TF_NEED_HDFS = true;
414 TF_ENABLE_XLA = tfFeature xlaSupport;
415
416 CC_OPT_FLAGS = " ";
417
418 # https://github.com/tensorflow/tensorflow/issues/14454
419 TF_NEED_MPI = tfFeature cudaSupport;
420
421 TF_NEED_CUDA = tfFeature cudaSupport;
422 TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkitDevMerged},${cudnnMerged},${lib.getLib nccl}";
423 TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
424
425 # Needed even when we override stdenv: e.g. for ar
426 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
427 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/cc";
428
429 patches = [
430 "${gentoo-patches}/0002-systemlib-Latest-absl-LTS-has-split-cord-libs.patch"
431 "${gentoo-patches}/0005-systemlib-Updates-for-Abseil-20220623-LTS.patch"
432 "${gentoo-patches}/0007-systemlibs-Add-well_known_types_py_pb2-target.patch"
433 # https://github.com/conda-forge/tensorflow-feedstock/pull/329/commits/0a63c5a962451b4da99a9948323d8b3ed462f461
434 (fetchpatch {
435 name = "fix-layout-proto-duplicate-loading.patch";
436 url = "https://raw.githubusercontent.com/conda-forge/tensorflow-feedstock/0a63c5a962451b4da99a9948323d8b3ed462f461/recipe/patches/0001-Omit-linking-to-layout_proto_cc-if-protobuf-linkage-.patch";
437 hash = "sha256-/7buV6DinKnrgfqbe7KKSh9rCebeQdXv2Uj+Xg/083w=";
438 })
439 ./com_google_absl_add_log.patch
440 ./absl_py_argparse_flags.patch
441 ./protobuf_python.patch
442 ./pybind11_protobuf_python_runtime_dep.patch
443 ./pybind11_protobuf_newer_version.patch
444 ] ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-darwin") [ ./absl_to_std.patch ];
445
446 postPatch =
447 ''
448 # bazel 3.3 should work just as well as bazel 3.1
449 rm -f .bazelversion
450 patchShebangs .
451 ''
452 + lib.optionalString (stdenv.hostPlatform.system == "x86_64-darwin") ''
453 cat ${./com_google_absl_fix_macos.patch} >> third_party/absl/com_google_absl_fix_mac_and_nvcc_build.patch
454 ''
455 + lib.optionalString (!withTensorboard) ''
456 # Tensorboard pulls in a bunch of dependencies, some of which may
457 # include security vulnerabilities. So we make it optional.
458 # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
459 sed -i '/tensorboard ~=/d' tensorflow/tools/pip_package/setup.py
460 '';
461
462 # https://github.com/tensorflow/tensorflow/pull/39470
463 env.NIX_CFLAGS_COMPILE = toString [ "-Wno-stringop-truncation" ];
464
465 preConfigure =
466 let
467 opt_flags =
468 [ ]
469 ++ lib.optionals sse42Support [ "-msse4.2" ]
470 ++ lib.optionals avx2Support [ "-mavx2" ]
471 ++ lib.optionals fmaSupport [ "-mfma" ];
472 in
473 ''
474 patchShebangs configure
475
476 # dummy ldconfig
477 mkdir dummy-ldconfig
478 echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
479 chmod +x dummy-ldconfig/ldconfig
480 export PATH="$PWD/dummy-ldconfig:$PATH"
481
482 export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
483 export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}"
484 mkdir -p "$PYTHON_LIB_PATH"
485
486 # To avoid mixing Python 2 and Python 3
487 unset PYTHONPATH
488 '';
489
490 configurePhase = ''
491 runHook preConfigure
492 ./configure
493 runHook postConfigure
494 '';
495
496 hardeningDisable = [ "format" ];
497
498 bazelBuildFlags =
499 [
500 "--config=opt" # optimize using the flags set in the configure phase
501 ]
502 ++ lib.optionals stdenv.cc.isClang [
503 "--cxxopt=-x"
504 "--cxxopt=c++"
505 "--host_cxxopt=-x"
506 "--host_cxxopt=c++"
507
508 # workaround for https://github.com/bazelbuild/bazel/issues/15359
509 "--spawn_strategy=sandboxed"
510 ]
511 ++ lib.optionals (mklSupport) [ "--config=mkl" ];
512
513 bazelTargets = [
514 "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow"
515 ];
516
517 removeRulesCC = false;
518 # Without this Bazel complaints about sandbox violations.
519 dontAddBazelOpts = true;
520
521 fetchAttrs = {
522 sha256 =
523 {
524 x86_64-linux =
525 if cudaSupport then
526 "sha256-5VFMNHeLrUxW5RTr6EhT3pay9nWJ5JkZTGirDds5QkU="
527 else
528 "sha256-KzgWV69Btr84FdwQ5JI2nQEsqiPg1/+TWdbw5bmxXOE=";
529 aarch64-linux =
530 if cudaSupport then
531 "sha256-ty5+51BwHWE1xR4/0WcWTp608NzSAS/iiyN+9zx7/wI="
532 else
533 "sha256-9btXrNHqd720oXTPDhSmFidv5iaZRLjCVX8opmrMjXk=";
534 x86_64-darwin = "sha256-gqb03kB0z2pZQ6m1fyRp1/Nbt8AVVHWpOJSeZNCLc4w=";
535 aarch64-darwin = "sha256-WdgAaFZU+ePwWkVBhLzjlNT7ELfGHOTaMdafcAMD5yo=";
536 }
537 .${stdenv.hostPlatform.system} or (throw "unsupported system ${stdenv.hostPlatform.system}");
538 };
539
540 buildAttrs = {
541 outputs = [
542 "out"
543 "python"
544 ];
545
546 # need to rebuild schemas since we use a different flatbuffers version
547 preBuild = ''
548 (cd tensorflow/lite/schema;${flatbuffers-core}/bin/flatc --gen-object-api -c schema.fbs)
549 (cd tensorflow/lite/schema;${flatbuffers-core}/bin/flatc --gen-object-api -c conversion_metadata.fbs)
550 (cd tensorflow/lite/acceleration/configuration;${flatbuffers-core}/bin/flatc -o configuration.fbs --proto configuration.proto)
551 sed -i s,tflite.proto,tflite,g tensorflow/lite/acceleration/configuration/configuration.fbs/configuration.fbs
552 (cd tensorflow/lite/acceleration/configuration;${flatbuffers-core}/bin/flatc --gen-compare --gen-object-api -c configuration.fbs/configuration.fbs)
553 cp -r tensorflow/lite/acceleration/configuration/configuration.fbs tensorflow/lite/experimental/acceleration/configuration
554 (cd tensorflow/lite/experimental/acceleration/configuration;${flatbuffers-core}/bin/flatc -c configuration.fbs/configuration.fbs)
555 (cd tensorflow/lite/delegates/gpu/cl;${flatbuffers-core}/bin/flatc -c compiled_program_cache.fbs)
556 (cd tensorflow/lite/delegates/gpu/cl;${flatbuffers-core}/bin/flatc -I $NIX_BUILD_TOP/source -c serialization.fbs)
557 (cd tensorflow/lite/delegates/gpu/common;${flatbuffers-core}/bin/flatc -I $NIX_BUILD_TOP/source -c gpu_model.fbs)
558 (cd tensorflow/lite/delegates/gpu/common/task;${flatbuffers-core}/bin/flatc -c serialization_base.fbs)
559 patchShebangs .
560 '';
561
562 installPhase = ''
563 mkdir -p "$out"
564 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
565 # Write pkgconfig file.
566 mkdir "$out/lib/pkgconfig"
567 cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
568 Name: TensorFlow
569 Version: ${version}
570 Description: Library for computation using data flow graphs for scalable machine learning
571 Requires:
572 Libs: -L$out/lib -ltensorflow
573 Cflags: -I$out/include/tensorflow
574 EOF
575
576 # build the source code, then copy it to $python (build_pip_package
577 # actually builds a symlink farm so we must dereference them).
578 bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
579 cp -Lr "$PWD/dist" "$python"
580 '';
581
582 postFixup = lib.optionalString cudaSupport ''
583 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
584 addDriverRunpath "$lib"
585 done
586 '';
587
588 requiredSystemFeatures = [ "big-parallel" ];
589 };
590
591 meta =
592 with lib;
593 {
594 badPlatforms = lib.optionals cudaSupport lib.platforms.darwin;
595 changelog = "https://github.com/tensorflow/tensorflow/releases/tag/v${version}";
596 description = "Computation using data flow graphs for scalable machine learning";
597 homepage = "http://tensorflow.org";
598 license = licenses.asl20;
599 maintainers = with maintainers; [ abbradar ];
600 platforms = with platforms; linux ++ darwin;
601 broken =
602 stdenv.isDarwin
603 || !(xlaSupport -> cudaSupport)
604 || !(cudaSupport -> builtins.hasAttr cudnnAttribute cudaPackages)
605 || !(cudaSupport -> cudaPackages ? cudatoolkit);
606 }
607 // lib.optionalAttrs stdenv.isDarwin {
608 timeout = 86400; # 24 hours
609 maxSilent = 14400; # 4h, double the default of 7200s
610 };
611 };
612in
613buildPythonPackage {
614 inherit version pname;
615 disabled = pythonAtLeast "3.12";
616
617 src = bazel-build.python;
618
619 # Adjust dependency requirements:
620 # - Drop tensorflow-io dependency until we get it to build
621 # - Relax flatbuffers and gast version requirements
622 # - The purpose of python3Packages.libclang is not clear at the moment and we don't have it packaged yet
623 # - keras and tensorlow-io-gcs-filesystem will be considered as optional for now.
624 postPatch = ''
625 sed -i setup.py \
626 -e '/tensorflow-io-gcs-filesystem/,+1d' \
627 -e "s/'flatbuffers[^']*',/'flatbuffers',/" \
628 -e "s/'gast[^']*',/'gast',/" \
629 -e "/'libclang[^']*',/d" \
630 -e "/'keras[^']*')\?,/d" \
631 -e "/'tensorflow-io-gcs-filesystem[^']*',/d" \
632 -e "s/'protobuf[^']*',/'protobuf',/" \
633 '';
634
635 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
636 # and the propagated input tensorboard, which causes environment collisions.
637 # Another possibility would be to have tensorboard only in the buildInputs
638 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
639 postInstall = ''
640 rm $out/bin/tensorboard
641 '';
642
643 setupPyGlobalFlags = [ "--project_name ${pname}" ];
644
645 # tensorflow/tools/pip_package/setup.py
646 propagatedBuildInputs = [
647 absl-py
648 abseil-cpp
649 astunparse
650 flatbuffers-python
651 gast
652 google-pasta
653 grpcio
654 h5py
655 keras-preprocessing
656 numpy
657 opt-einsum
658 packaging
659 protobuf-python
660 six
661 tensorflow-estimator-bin
662 termcolor
663 typing-extensions
664 wrapt
665 ] ++ lib.optionals withTensorboard [ tensorboard ];
666
667 nativeBuildInputs = lib.optionals cudaSupport [ addDriverRunpath ];
668
669 postFixup = lib.optionalString cudaSupport ''
670 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
671 addDriverRunpath "$lib"
672
673 patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:${cudnnMerged}/lib:${lib.getLib nccl}/lib:$(patchelf --print-rpath "$lib")" "$lib"
674 done
675 '';
676
677 # Actual tests are slow and impure.
678 # TODO try to run them anyway
679 # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
680 # TEST_PACKAGES in tensorflow/tools/pip_package/setup.py
681 nativeCheckInputs = [
682 dill
683 portpicker
684 tblib
685 ];
686 checkPhase = ''
687 ${python.interpreter} <<EOF
688 # A simple "Hello world"
689 import tensorflow as tf
690 hello = tf.constant("Hello, world!")
691 tf.print(hello)
692
693 tf.random.set_seed(0)
694 width = 512
695 choice = 48
696 t_in = tf.Variable(tf.random.uniform(shape=[width]))
697 with tf.GradientTape() as tape:
698 t_out = tf.slice(tf.nn.softmax(t_in), [choice], [1])
699 diff = tape.gradient(t_out, t_in)
700 assert(0 < tf.reduce_min(tf.slice(diff, [choice], [1])))
701 assert(0 > tf.reduce_max(tf.slice(diff, [1], [choice - 1])))
702 EOF
703 '';
704 # Regression test for #77626 removed because not more `tensorflow.contrib`.
705
706 passthru = {
707 deps = bazel-build.deps;
708 libtensorflow = bazel-build.out;
709 };
710
711 inherit (bazel-build) meta;
712}