tjh.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 cmake,
8 pybind11,
9 setuptools,
10
11 # buildInputs
12 abseil-cpp,
13 protobuf,
14 gtest,
15
16 # dependencies
17 numpy,
18 typing-extensions,
19
20 # tests
21 google-re2,
22 ml-dtypes,
23 nbval,
24 parameterized,
25 pillow,
26 pytestCheckHook,
27 tabulate,
28 writableTmpDirAsHomeHook,
29}:
30
31let
32 gtestStatic = gtest.override { static = true; };
33in
34buildPythonPackage rec {
35 pname = "onnx";
36 version = "1.18.0";
37 pyproject = true;
38
39 src = fetchFromGitHub {
40 owner = "onnx";
41 repo = "onnx";
42 tag = "v${version}";
43 hash = "sha256-UhtF+CWuyv5/Pq/5agLL4Y95YNP63W2BraprhRqJOag=";
44 };
45
46 build-system = [
47 cmake
48 protobuf
49 setuptools
50 ];
51
52 buildInputs = [
53 abseil-cpp
54 gtestStatic
55 pybind11
56 ];
57
58 dependencies = [
59 protobuf
60 numpy
61 typing-extensions
62 ];
63
64 nativeCheckInputs = [
65 google-re2
66 ml-dtypes
67 nbval
68 parameterized
69 pillow
70 pytestCheckHook
71 tabulate
72 writableTmpDirAsHomeHook
73 ];
74
75 postPatch = ''
76 rm -r third_party
77
78 chmod +x tools/protoc-gen-mypy.sh.in
79 patchShebangs tools/protoc-gen-mypy.sh.in
80 '';
81
82 preConfigure = ''
83 # Set CMAKE_INSTALL_LIBDIR to lib explicitly, because otherwise it gets set
84 # to lib64 and cmake incorrectly looks for the protobuf library in lib64
85 export CMAKE_ARGS="-DCMAKE_INSTALL_LIBDIR=lib -DONNX_USE_PROTOBUF_SHARED_LIBS=ON"
86 export CMAKE_ARGS+=" -Dgoogletest_STATIC_LIBRARIES=${gtestStatic}/lib/libgtest.a"
87 export ONNX_BUILD_TESTS=1
88 '';
89
90 preBuild = ''
91 export MAX_JOBS=$NIX_BUILD_CORES
92 '';
93
94 # The executables are just utility scripts that aren't too important
95 postInstall = ''
96 rm -r $out/bin
97 '';
98
99 # The setup.py does all the configuration
100 dontUseCmakeConfigure = true;
101
102 # detecting source dir as a python package confuses pytest
103 preCheck = ''
104 rm onnx/__init__.py
105 '';
106
107 enabledTestPaths = [
108 "onnx/test"
109 "examples"
110 ];
111
112 __darwinAllowLocalNetworking = true;
113
114 postCheck = ''
115 # run "cpp" tests
116 .setuptools-cmake-build/onnx_gtests
117 '';
118
119 pythonImportsCheck = [ "onnx" ];
120
121 meta = {
122 description = "Open Neural Network Exchange";
123 homepage = "https://onnx.ai";
124 changelog = "https://github.com/onnx/onnx/releases/tag/v${version}";
125 license = lib.licenses.asl20;
126 maintainers = with lib.maintainers; [ acairncross ];
127 };
128}