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