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