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