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