pyrox.dev
lol
1{ stdenv
2, bash
3, buildPackages
4, buildBazelPackage
5, fetchFromGitHub
6, lib
7}:
8let
9 buildPlatform = stdenv.buildPlatform;
10 hostPlatform = stdenv.hostPlatform;
11 pythonEnv = buildPackages.python3.withPackages (ps: [ ps.numpy ]);
12 bazelDepsSha256ByBuildAndHost = {
13 x86_64-linux = {
14 x86_64-linux = "sha256-61qmnAB80syYhURWYJOiOnoGOtNa1pPkxfznrFScPAo=";
15 aarch64-linux = "sha256-sOIYpp98wJRz3RGvPasyNEJ05W29913Lsm+oi/aq/Ag=";
16 };
17 aarch64-linux = {
18 aarch64-linux = "sha256-MJU4y9Dt9xJWKgw7iKW+9Ur856rMIHeFD5u05s+Q7rQ=";
19 };
20 };
21 bazelHostConfigName.aarch64-linux = "elinux_aarch64";
22 bazelDepsSha256ByHost =
23 bazelDepsSha256ByBuildAndHost.${buildPlatform.system} or
24 (throw "unsupported build system ${buildPlatform.system}");
25 bazelDepsSha256 = bazelDepsSha256ByHost.${hostPlatform.system} or
26 (throw "unsupported host system ${hostPlatform.system} with build system ${buildPlatform.system}");
27in
28buildBazelPackage rec {
29 name = "tensorflow-lite";
30 version = "2.13.0";
31
32 src = fetchFromGitHub {
33 owner = "tensorflow";
34 repo = "tensorflow";
35 rev = "v${version}";
36 hash = "sha256-Rq5pAVmxlWBVnph20fkAwbfy+iuBNlfFy14poDPd5h0=";
37 };
38
39 bazel = buildPackages.bazel_5;
40
41 nativeBuildInputs = [ pythonEnv buildPackages.perl ];
42
43 bazelTargets = [
44 "//tensorflow/lite:libtensorflowlite.so"
45 "//tensorflow/lite/c:tensorflowlite_c"
46 "//tensorflow/lite/tools/benchmark:benchmark_model"
47 "//tensorflow/lite/tools/benchmark:benchmark_model_performance_options"
48 ];
49
50 bazelFlags = [
51 "--config=opt"
52 ] ++ lib.optionals (hostPlatform.system != buildPlatform.system) [
53 "--config=${bazelHostConfigName.${hostPlatform.system}}"
54 ];
55
56 bazelBuildFlags = [ "--cxxopt=--std=c++17" ];
57
58 buildAttrs = {
59 installPhase = ''
60 mkdir -p $out/{bin,lib}
61
62 # copy the libs and binaries into the output dir
63 cp ./bazel-bin/tensorflow/lite/c/libtensorflowlite_c.so $out/lib
64 cp ./bazel-bin/tensorflow/lite/libtensorflowlite.so $out/lib
65 cp ./bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model $out/bin
66 cp ./bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model_performance_options $out/bin
67
68 find . -type f -name '*.h' | while read f; do
69 path="$out/include/''${f/.\//}"
70 install -D "$f" "$path"
71
72 # remove executable bit from headers
73 chmod -x "$path"
74 done
75 '';
76 };
77
78 fetchAttrs.sha256 = bazelDepsSha256;
79
80 PYTHON_BIN_PATH = pythonEnv.interpreter;
81
82 dontAddBazelOpts = true;
83 removeRulesCC = false;
84
85 postPatch = ''
86 rm .bazelversion
87
88 # Fix gcc-13 build failure by including missing include headers
89 sed -e '1i #include <cstdint>' -i \
90 tensorflow/lite/kernels/internal/spectrogram.cc
91 '';
92
93 preConfigure = ''
94 patchShebangs configure
95 '';
96
97 # configure script freaks out when parameters are passed
98 dontAddPrefix = true;
99 configurePlatforms = [];
100
101 meta = with lib; {
102 description = "An open source deep learning framework for on-device inference.";
103 homepage = "https://www.tensorflow.org/lite";
104 license = licenses.asl20;
105 maintainers = with maintainers; [ mschwaig cpcloud ];
106 platforms = [ "x86_64-linux" "aarch64-linux" ];
107 };
108}