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