tjh.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 cmake,
6 fetchFromGitHub,
7 gtest,
8 nbval,
9 numpy,
10 parameterized,
11 protobuf_21,
12 pybind11,
13 pytestCheckHook,
14 pythonOlder,
15 tabulate,
16 typing-extensions,
17 abseil-cpp,
18 google-re2,
19 pillow,
20 protobuf,
21}:
22
23let
24 gtestStatic = gtest.override { static = true; };
25in
26buildPythonPackage rec {
27 pname = "onnx";
28 version = "1.16.1";
29 format = "setuptools";
30
31 disabled = pythonOlder "3.8";
32
33 src = fetchFromGitHub {
34 owner = pname;
35 repo = pname;
36 rev = "refs/tags/v${version}";
37 hash = "sha256-I1wwfn91hdH3jORIKny0Xc73qW2P04MjkVCgcaNnQUE=";
38 };
39
40 build-system = [
41 cmake
42 pybind11
43 ];
44
45 buildInputs = [
46 abseil-cpp
47 protobuf
48 google-re2
49 gtestStatic
50 pillow
51 ];
52
53 dependencies = [
54 protobuf_21
55 protobuf
56 numpy
57 typing-extensions
58 ];
59
60 nativeCheckInputs = [
61 nbval
62 parameterized
63 pytestCheckHook
64 tabulate
65 ];
66
67 postPatch = ''
68 chmod +x tools/protoc-gen-mypy.sh.in
69 patchShebangs tools/protoc-gen-mypy.sh.in
70
71 substituteInPlace setup.py \
72 --replace 'setup_requires.append("pytest-runner")' ""
73
74 # prevent from fetching & building own gtest
75 substituteInPlace CMakeLists.txt \
76 --replace 'include(googletest)' ""
77 substituteInPlace cmake/unittest.cmake \
78 --replace 'googletest)' ')'
79 '';
80
81 preConfigure = ''
82 # Set CMAKE_INSTALL_LIBDIR to lib explicitly, because otherwise it gets set
83 # to lib64 and cmake incorrectly looks for the protobuf library in lib64
84 export CMAKE_ARGS="-DCMAKE_INSTALL_LIBDIR=lib -DONNX_USE_PROTOBUF_SHARED_LIBS=ON"
85 export CMAKE_ARGS+=" -Dgoogletest_STATIC_LIBRARIES=${gtestStatic}/lib/libgtest.a -Dgoogletest_INCLUDE_DIRS=${lib.getDev gtestStatic}/include"
86 export ONNX_BUILD_TESTS=1
87 '';
88
89 preBuild = ''
90 export MAX_JOBS=$NIX_BUILD_CORES
91 '';
92
93 # The executables are just utility scripts that aren't too important
94 postInstall = ''
95 rm -r $out/bin
96 '';
97
98 # The setup.py does all the configuration
99 dontUseCmakeConfigure = true;
100
101 preCheck = ''
102 export HOME=$(mktemp -d)
103
104 # detecting source dir as a python package confuses pytest
105 mv onnx/__init__.py onnx/__init__.py.hidden
106 '';
107
108 pytestFlagsArray = [
109 "onnx/test"
110 "onnx/examples"
111 ];
112
113 disabledTests =
114 [
115 # attempts to fetch data from web
116 "test_bvlc_alexnet_cpu"
117 "test_densenet121_cpu"
118 "test_inception_v1_cpu"
119 "test_inception_v2_cpu"
120 "test_resnet50_cpu"
121 "test_shufflenet_cpu"
122 "test_squeezenet_cpu"
123 "test_vgg19_cpu"
124 "test_zfnet512_cpu"
125 ]
126 ++ lib.optionals stdenv.isAarch64 [
127 # AssertionError: Output 0 of test 0 in folder
128 "test__pytorch_converted_Conv2d_depthwise_padded"
129 "test__pytorch_converted_Conv2d_dilated"
130 "test_dft"
131 "test_dft_axis"
132 # AssertionError: Mismatch in test 'test_Conv2d_depthwise_padded'
133 "test_xor_bcast4v4d"
134 # AssertionError: assert 1 == 0
135 "test_ops_tested"
136 ];
137
138 disabledTestPaths = [
139 # Unexpected output fields from running code: {'stderr'}
140 "onnx/examples/np_array_tensorproto.ipynb"
141 ];
142
143 __darwinAllowLocalNetworking = true;
144
145 postCheck = ''
146 # run "cpp" tests
147 .setuptools-cmake-build/onnx_gtests
148 '';
149
150 pythonImportsCheck = [ "onnx" ];
151
152 meta = with lib; {
153 description = "Open Neural Network Exchange";
154 homepage = "https://onnx.ai";
155 license = licenses.asl20;
156 maintainers = with maintainers; [ acairncross ];
157 };
158}