gen_onnx.py 43.2 KB
Newer Older
Khalique's avatar
Khalique committed
1
2
3
4
5
6
import numpy as np
import onnx
from onnx import helper
from onnx import numpy_helper
from onnx import AttributeProto, TensorProto, GraphProto

Khalique's avatar
Khalique committed
7

Khalique's avatar
Khalique committed
8
9
def onnx_test(op_test):
    def run_test():
Khalique's avatar
Khalique committed
10
11
        op_info = op_test()
        if len(op_info) > 3:
Khalique's avatar
Khalique committed
12
13
14
15
16
            graph_def = helper.make_graph(op_info[0],
                                          op_test.__name__,
                                          op_info[1],
                                          op_info[2],
                                          initializer=op_info[3])
Khalique's avatar
Khalique committed
17
        else:
Khalique's avatar
Khalique committed
18
19
20
21
            graph_def = helper.make_graph(op_info[0], op_test.__name__,
                                          op_info[1], op_info[2])
        model_def = helper.make_model(graph_def,
                                      producer_name=op_test.__name__)
Khalique's avatar
Khalique committed
22
        onnx.save(model_def, '{}.onnx'.format(op_test.__name__))
Khalique's avatar
Khalique committed
23

Khalique's avatar
Khalique committed
24
25
    return run_test

Khalique's avatar
Khalique committed
26

Khalique's avatar
Khalique committed
27
@onnx_test
Khalique's avatar
Khalique committed
28
29
30
31
32
33
34
35
36
37
def acos_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Acos',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
38
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
39

Khalique's avatar
Khalique committed
40

Khalique's avatar
Khalique committed
41
@onnx_test
Khalique's avatar
Khalique committed
42
43
44
def add_bcast_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 4])
Khalique's avatar
Khalique committed
45
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [2, 3, 4, 5])
Khalique's avatar
Khalique committed
46

Khalique's avatar
Khalique committed
47
48
49
50
51
52
53
    node = onnx.helper.make_node('Add',
                                 inputs=['0', '1'],
                                 broadcast=1,
                                 axis=1,
                                 outputs=['2'])

    return ([node], [x, y], [z])
Khalique's avatar
Khalique committed
54
55


Khalique's avatar
Khalique committed
56
@onnx_test
Khalique's avatar
Khalique committed
57
58
59
60
61
62
63
64
65
66
67
def add_fp16_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT16, [1])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT16, [1])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT16, [1])

    node = onnx.helper.make_node(
        'Add',
        inputs=['0', '1'],
        outputs=['2'],
    )

Khalique's avatar
Khalique committed
68
    return (
Khalique's avatar
Khalique committed
69
        [node],
Khalique's avatar
Khalique committed
70
        [x, y],
Khalique's avatar
Khalique committed
71
72
        [z],
        # '0' -> 1.5, '1' -> 2.5
Khalique's avatar
Khalique committed
73
74
75
76
        [
            onnx.helper.make_tensor('0', TensorProto.FLOAT16, [1], [15872]),
            onnx.helper.make_tensor('1', TensorProto.FLOAT16, [1], [16640])
        ])
Khalique's avatar
Khalique committed
77
78


Khalique's avatar
Khalique committed
79
@onnx_test
Khalique's avatar
Khalique committed
80
81
82
def add_scalar_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [])
Khalique's avatar
Khalique committed
83
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [2, 3, 4, 5])
Khalique's avatar
Khalique committed
84

Khalique's avatar
Khalique committed
85
86
87
88
    node = onnx.helper.make_node('Add', inputs=['0', '1'], outputs=['2'])

    return ([node], [x, y], [z],
            [helper.make_tensor('1', TensorProto.FLOAT, [], [1])])
Khalique's avatar
Khalique committed
89
90


Khalique's avatar
Khalique committed
91
@onnx_test
Khalique's avatar
Khalique committed
92
93
94
95
def argmax_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 6])

Khalique's avatar
Khalique committed
96
97
98
99
100
    node = onnx.helper.make_node('ArgMax',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axis=2,
                                 keepdims=0)
Khalique's avatar
Khalique committed
101

Khalique's avatar
Khalique committed
102
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
103

Khalique's avatar
Khalique committed
104

Khalique's avatar
Khalique committed
105
@onnx_test
Khalique's avatar
Khalique committed
106
107
108
109
def argmin_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 5])

Khalique's avatar
Khalique committed
110
111
112
113
114
    node = onnx.helper.make_node('ArgMin',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axis=3,
                                 keepdims=0)
Khalique's avatar
Khalique committed
115

Khalique's avatar
Khalique committed
116
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
117

Khalique's avatar
Khalique committed
118

Khalique's avatar
Khalique committed
119
@onnx_test
Khalique's avatar
Khalique committed
120
121
122
123
124
125
126
127
128
129
def asin_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Asin',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
130
131
    return ([node], [x], [y])

Khalique's avatar
Khalique committed
132

Khalique's avatar
Khalique committed
133
@onnx_test
Khalique's avatar
Khalique committed
134
135
136
137
138
139
140
141
142
def atan_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Atan',
        inputs=['x'],
        outputs=['y'],
    )
Khalique's avatar
Khalique committed
143

Khalique's avatar
Khalique committed
144
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
145

Khalique's avatar
Khalique committed
146

Khalique's avatar
Khalique committed
147
@onnx_test
Khalique's avatar
Khalique committed
148
149
150
151
def cast_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT16, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

Khalique's avatar
Khalique committed
152
153
    node = onnx.helper.make_node('Cast', inputs=['x'], outputs=['y'], to=1)

Khalique's avatar
Khalique committed
154
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
155

Khalique's avatar
Khalique committed
156

Khalique's avatar
Khalique committed
157
@onnx_test
Khalique's avatar
Khalique committed
158
159
160
161
def clip_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])

Khalique's avatar
Khalique committed
162
163
164
165
166
    node = onnx.helper.make_node('Clip',
                                 inputs=['0'],
                                 outputs=['1'],
                                 max=6.0,
                                 min=0.0)
Khalique's avatar
Khalique committed
167

Khalique's avatar
Khalique committed
168
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
169

Khalique's avatar
Khalique committed
170

Khalique's avatar
Khalique committed
171
@onnx_test
Khalique's avatar
Khalique committed
172
173
174
175
176
177
178
179
180
181
182
183
def concat_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 4, 3])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [7, 4, 3])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [9, 4, 3])

    node = onnx.helper.make_node(
        'Concat',
        inputs=['0', '1'],
        axis=0,
        outputs=['2'],
    )

Khalique's avatar
Khalique committed
184
185
    return ([node], [x, y], [z])

Khalique's avatar
Khalique committed
186

Khalique's avatar
Khalique committed
187
@onnx_test
Khalique's avatar
Khalique committed
188
189
190
def constant_test():
    x = np.array([0, 1, 2])
    y = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
Khalique's avatar
Khalique committed
191

Khalique's avatar
Khalique committed
192
193
194
195
196
197
198
199
200
201
202
203
    node = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['0'],
        value=onnx.helper.make_tensor(
            name='const_tensor',
            data_type=TensorProto.FLOAT,
            dims=x.shape,
            vals=x.flatten().astype(float),
        ),
    )

Khalique's avatar
Khalique committed
204
    return ([node], [], [y])
Khalique's avatar
Khalique committed
205

Khalique's avatar
Khalique committed
206

Khalique's avatar
Khalique committed
207
@onnx_test
Khalique's avatar
Khalique committed
208
def constant_fill_test():
Khalique's avatar
Khalique committed
209
210
211
212
213
214
    value = helper.make_tensor_value_info('value', TensorProto.FLOAT, [2, 3])

    node = onnx.helper.make_node(
        'ConstantFill',
        inputs=[],
        outputs=['value'],
Khalique's avatar
Khalique committed
215
216
217
218
        dtype=1,
        value=1.0,
        shape=[2, 3],
        input_as_shape=0,
Khalique's avatar
Khalique committed
219
220
    )

Khalique's avatar
Khalique committed
221
    return ([node], [], [value])
Khalique's avatar
Khalique committed
222

Khalique's avatar
Khalique committed
223

Khalique's avatar
Khalique committed
224
@onnx_test
Khalique's avatar
Khalique committed
225
def constant_fill_input_as_shape_test():
Khalique's avatar
Khalique committed
226
    np_shape = np.array([2, 3])
Khalique's avatar
Khalique committed
227
228
229
    shape = helper.make_tensor_value_info('shape', TensorProto.INT32, [2])
    value = helper.make_tensor_value_info('value', TensorProto.FLOAT, [2, 3])

Khalique's avatar
Khalique committed
230
231
232
233
    ts_shape = helper.make_tensor(name='shape_tensor',
                                  data_type=TensorProto.INT32,
                                  dims=np_shape.shape,
                                  vals=np_shape.flatten().astype(int))
Khalique's avatar
Khalique committed
234
235
236
237
238
239
240
241
242
243
244
245

    const_shape_node = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['shape'],
        value=ts_shape,
    )

    node = onnx.helper.make_node(
        'ConstantFill',
        inputs=['shape'],
        outputs=['value'],
Khalique's avatar
Khalique committed
246
247
248
        dtype=1,
        value=1.0,
        input_as_shape=1,
Khalique's avatar
Khalique committed
249
250
    )

Khalique's avatar
Khalique committed
251
    return ([const_shape_node, node], [], [value])
Khalique's avatar
Khalique committed
252

Khalique's avatar
Khalique committed
253

Khalique's avatar
Khalique committed
254
@onnx_test
Khalique's avatar
Khalique committed
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def constant_scalar_test():
    x = np.array([1])
    y = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1])

    node = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['0'],
        value=onnx.helper.make_tensor(
            name='const_tensor',
            data_type=TensorProto.INT32,
            dims=x.shape,
            vals=x.flatten().astype(int),
        ),
    )

Khalique's avatar
Khalique committed
271
    return ([node], [], [y])
Khalique's avatar
Khalique committed
272

Khalique's avatar
Khalique committed
273

Khalique's avatar
Khalique committed
274
@onnx_test
Khalique's avatar
Khalique committed
275
def const_of_shape_empty_input_test():
Khalique's avatar
Khalique committed
276
277
    tensor_val = onnx.helper.make_tensor('value', onnx.TensorProto.INT64, [1],
                                         [10])
Khalique's avatar
Khalique committed
278
279
    shape_val = np.array([2, 3, 4]).astype(np.int64)
    empty_val = np.array([]).astype(np.int64)
Khalique's avatar
Khalique committed
280
281
282
283
    empty_ts = helper.make_tensor(name='empty_tensor',
                                  data_type=TensorProto.INT32,
                                  dims=empty_val.shape,
                                  vals=empty_val.flatten().astype(int))
Khalique's avatar
Khalique committed
284
285
286
287
288
289
290
291
292
293
294
295
    shape_const = helper.make_node(
        'Constant',
        inputs=[],
        outputs=['shape'],
        value=empty_ts,
    )
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 3, 4])

    node = onnx.helper.make_node(
        'ConstantOfShape',
        inputs=['shape'],
        outputs=['y'],
Khalique's avatar
Khalique committed
296
        value=tensor_val,
Khalique's avatar
Khalique committed
297
298
    )

Khalique's avatar
Khalique committed
299
    return ([shape_const, node], [], [y])
Khalique's avatar
Khalique committed
300

Khalique's avatar
Khalique committed
301

Khalique's avatar
Khalique committed
302
@onnx_test
Khalique's avatar
Khalique committed
303
def const_of_shape_float_test():
Khalique's avatar
Khalique committed
304
305
    tensor_val = onnx.helper.make_tensor('value', onnx.TensorProto.FLOAT, [1],
                                         [10])
Khalique's avatar
Khalique committed
306
307

    shape_val = np.array([2, 3, 4]).astype(np.int64)
Khalique's avatar
Khalique committed
308
309
310
311
    shape_ts = helper.make_tensor(name='shape_tensor',
                                  data_type=TensorProto.INT32,
                                  dims=shape_val.shape,
                                  vals=shape_val.flatten().astype(int))
Khalique's avatar
Khalique committed
312
313
314
315
316
317
318
319
320

    shape_const = helper.make_node(
        'Constant',
        inputs=[],
        outputs=['shape'],
        value=shape_ts,
    )
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 3, 4])

Khalique's avatar
Khalique committed
321
322
323
324
    node = onnx.helper.make_node('ConstantOfShape',
                                 inputs=['shape'],
                                 outputs=['y'],
                                 value=tensor_val)
Khalique's avatar
Khalique committed
325

Khalique's avatar
Khalique committed
326
    return ([shape_const, node], [], [y])
Khalique's avatar
Khalique committed
327

Khalique's avatar
Khalique committed
328

Khalique's avatar
Khalique committed
329
@onnx_test
Khalique's avatar
Khalique committed
330
def const_of_shape_int64_test():
Khalique's avatar
Khalique committed
331
332
    tensor_val = onnx.helper.make_tensor('value', onnx.TensorProto.INT64, [1],
                                         [10])
Khalique's avatar
Khalique committed
333
    shape_val = np.array([2, 3, 4]).astype(np.int64)
Khalique's avatar
Khalique committed
334
335
336
337
    shape_ts = helper.make_tensor(name='shape_tensor',
                                  data_type=TensorProto.INT32,
                                  dims=shape_val.shape,
                                  vals=shape_val.flatten().astype(int))
Khalique's avatar
Khalique committed
338
    shape_const = helper.make_node(
Khalique's avatar
Khalique committed
339
340
341
342
        'Constant',
        inputs=[],
        outputs=['shape'],
        value=shape_ts,
Khalique's avatar
Khalique committed
343
344
    )
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 3, 4])
Khalique's avatar
Khalique committed
345
346
347
348
349

    node = onnx.helper.make_node('ConstantOfShape',
                                 inputs=['shape'],
                                 outputs=['y'],
                                 value=tensor_val)
Khalique's avatar
Khalique committed
350

Khalique's avatar
Khalique committed
351
    return ([shape_const, node], [], [y])
Khalique's avatar
Khalique committed
352

Khalique's avatar
Khalique committed
353

Khalique's avatar
Khalique committed
354
@onnx_test
Khalique's avatar
Khalique committed
355
356
def const_of_shape_no_value_attr_test():
    shape_val = np.array([2, 3, 4]).astype(np.int64)
Khalique's avatar
Khalique committed
357
358
359
360
    shape_ts = helper.make_tensor(name='shape_tensor',
                                  data_type=TensorProto.INT32,
                                  dims=shape_val.shape,
                                  vals=shape_val.flatten().astype(int))
Khalique's avatar
Khalique committed
361
362
363
364
365
366
367
    shape_const = helper.make_node(
        'Constant',
        inputs=[],
        outputs=['shape'],
        value=shape_ts,
    )
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 3, 4])
Khalique's avatar
Khalique committed
368

Khalique's avatar
Khalique committed
369
370
371
372
373
374
    node = onnx.helper.make_node(
        'ConstantOfShape',
        inputs=['shape'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
375
    return ([shape_const, node], [], [y])
Khalique's avatar
Khalique committed
376

Khalique's avatar
Khalique committed
377

Khalique's avatar
Khalique committed
378
@onnx_test
Khalique's avatar
Khalique committed
379
380
381
382
383
def conv_autopad_fail_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 32, 32])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 1, 1])
    out = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1, 1, 34, 34])

Khalique's avatar
Khalique committed
384
385
386
387
388
389
390
391
392
    node = onnx.helper.make_node('Conv',
                                 inputs=['0', '1'],
                                 outputs=['2'],
                                 dilations=[1, 1],
                                 strides=[1, 1],
                                 auto_pad='SAME',
                                 pads=[0, 0, 1, 1, 0, 0, 1, 1])

    return ([node], [x, y], [out])
Khalique's avatar
Khalique committed
393
394


Khalique's avatar
Khalique committed
395
@onnx_test
Khalique's avatar
Khalique committed
396
397
398
399
400
401
def conv_bias_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 32, 32])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 5, 5])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1])
    out = helper.make_tensor_value_info('3', TensorProto.FLOAT, [1, 2, 28, 28])

Khalique's avatar
Khalique committed
402
403
404
405
406
407
408
    node = onnx.helper.make_node('Conv',
                                 inputs=['0', '1', '2'],
                                 outputs=['3'],
                                 dilations=[1, 1],
                                 strides=[1, 1])

    return ([node], [x, y, z], [out])
Khalique's avatar
Khalique committed
409
410


Khalique's avatar
Khalique committed
411
@onnx_test
Khalique's avatar
Khalique committed
412
413
414
415
416
417
418
419
def conv_bn_relu_maxpool_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 32, 32])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 5, 5])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1])
    m = helper.make_tensor_value_info('3', TensorProto.FLOAT, [1])
    n = helper.make_tensor_value_info('4', TensorProto.FLOAT, [1])
    k = helper.make_tensor_value_info('5', TensorProto.FLOAT, [1])
    l = helper.make_tensor_value_info('6', TensorProto.FLOAT, [1])
Khalique's avatar
Khalique committed
420
421
    out = helper.make_tensor_value_info('10', TensorProto.FLOAT,
                                        [1, 1, 14, 14])
Khalique's avatar
Khalique committed
422

Khalique's avatar
Khalique committed
423
424
425
426
427
428
    node0 = onnx.helper.make_node('Conv',
                                  inputs=['0', '1', '2'],
                                  outputs=['7'],
                                  dilations=[1, 1],
                                  strides=[1, 1],
                                  pads=[0, 0, 0, 0])
Khalique's avatar
Khalique committed
429

Khalique's avatar
Khalique committed
430
431
432
433
434
    node1 = onnx.helper.make_node('BatchNormalization',
                                  inputs=['7', '3', '4', '5', '6'],
                                  outputs=['8'],
                                  epsilon=9.99999974737875e-06,
                                  momentum=0.899999976158142)
Khalique's avatar
Khalique committed
435

Khalique's avatar
Khalique committed
436
437
438
439
440
441
442
443
444
    node2 = onnx.helper.make_node('Relu', inputs=['8'], outputs=['9'])
    node3 = onnx.helper.make_node('MaxPool',
                                  inputs=['9'],
                                  outputs=['10'],
                                  pads=[0, 0, 0, 0],
                                  strides=[2, 2],
                                  kernel_shape=[2, 2])

    return ([node0, node1, node2, node3], [x, y, z, m, n, k, l], [out])
Khalique's avatar
Khalique committed
445
446


Khalique's avatar
Khalique committed
447
@onnx_test
Khalique's avatar
Khalique committed
448
449
450
451
452
453
def conv_relu_maxpool_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 32, 32])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 5, 5])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1])
    out = helper.make_tensor_value_info('5', TensorProto.FLOAT, [1, 1, 14, 14])

Khalique's avatar
Khalique committed
454
455
456
457
458
459
    node1 = onnx.helper.make_node('Conv',
                                  inputs=['0', '1', '2'],
                                  outputs=['3'],
                                  dilations=[1, 1],
                                  strides=[1, 1],
                                  pads=[0, 0, 0, 0])
Khalique's avatar
Khalique committed
460

Khalique's avatar
Khalique committed
461
    node2 = onnx.helper.make_node('Relu', inputs=['3'], outputs=['4'])
Khalique's avatar
Khalique committed
462

Khalique's avatar
Khalique committed
463
464
465
466
467
468
469
470
    node3 = onnx.helper.make_node('MaxPool',
                                  inputs=['4'],
                                  outputs=['5'],
                                  pads=[0, 0, 0, 0],
                                  strides=[2, 2],
                                  kernel_shape=[2, 2])

    return ([node1, node2, node3], [x, y, z], [out])
Khalique's avatar
Khalique committed
471
472


Khalique's avatar
Khalique committed
473
@onnx_test
Khalique's avatar
Khalique committed
474
475
476
477
478
479
480
481
def conv_relu_maxpool_x2_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 32, 32])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [5, 3, 5, 5])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [5])
    m = helper.make_tensor_value_info('3', TensorProto.FLOAT, [1, 5, 5, 5])
    n = helper.make_tensor_value_info('4', TensorProto.FLOAT, [1])
    out = helper.make_tensor_value_info('10', TensorProto.FLOAT, [1, 1, 5, 5])

Khalique's avatar
Khalique committed
482
483
484
485
486
487
    node1 = onnx.helper.make_node('Conv',
                                  inputs=['0', '1', '2'],
                                  outputs=['5'],
                                  dilations=[1, 1],
                                  strides=[1, 1],
                                  pads=[0, 0, 0, 0])
Khalique's avatar
Khalique committed
488

Khalique's avatar
Khalique committed
489
    node2 = onnx.helper.make_node('Relu', inputs=['5'], outputs=['6'])
Khalique's avatar
Khalique committed
490

Khalique's avatar
Khalique committed
491
492
493
494
495
496
    node3 = onnx.helper.make_node('MaxPool',
                                  inputs=['6'],
                                  outputs=['7'],
                                  pads=[0, 0, 0, 0],
                                  strides=[2, 2],
                                  kernel_shape=[2, 2])
Khalique's avatar
Khalique committed
497

Khalique's avatar
Khalique committed
498
499
500
501
502
503
    node4 = onnx.helper.make_node('Conv',
                                  inputs=['7', '3', '4'],
                                  outputs=['8'],
                                  dilations=[1, 1],
                                  strides=[1, 1],
                                  pads=[0, 0, 0, 0])
Khalique's avatar
Khalique committed
504

Khalique's avatar
Khalique committed
505
    node5 = onnx.helper.make_node('Relu', inputs=['8'], outputs=['9'])
Khalique's avatar
Khalique committed
506

Khalique's avatar
Khalique committed
507
508
509
510
511
512
513
514
    node6 = onnx.helper.make_node('MaxPool',
                                  inputs=['9'],
                                  outputs=['10'],
                                  pads=[0, 0, 0, 0],
                                  strides=[2, 2],
                                  kernel_shape=[2, 2])

    return ([node1, node2, node3, node4, node5, node6], [x, y, z, m, n], [out])
Khalique's avatar
Khalique committed
515
516


Khalique's avatar
Khalique committed
517
@onnx_test
Khalique's avatar
Khalique committed
518
519
520
521
522
523
524
525
526
527
def cos_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Cos',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
528
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
529

Khalique's avatar
Khalique committed
530

Khalique's avatar
Khalique committed
531
@onnx_test
Khalique's avatar
Khalique committed
532
533
534
535
536
537
538
539
540
541
def cosh_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1])

    node = onnx.helper.make_node(
        'Cosh',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
542
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
543

Khalique's avatar
Khalique committed
544

Khalique's avatar
Khalique committed
545
@onnx_test
Khalique's avatar
Khalique committed
546
def dropout_test():
Khalique's avatar
Khalique committed
547
548
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 2, 2])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 2, 2])
Khalique's avatar
Khalique committed
549

Khalique's avatar
Khalique committed
550
551
552
553
554
555
556
    node = onnx.helper.make_node(
        'Dropout',
        inputs=['0'],
        outputs=['1'],
    )

    return ([node], [x], [y])
Khalique's avatar
Khalique committed
557
558


Khalique's avatar
Khalique committed
559
@onnx_test
Khalique's avatar
Khalique committed
560
561
562
563
def elu_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])

Khalique's avatar
Khalique committed
564
565
566
567
    node = onnx.helper.make_node('Elu',
                                 inputs=['0'],
                                 outputs=['1'],
                                 alpha=0.01)
Khalique's avatar
Khalique committed
568

Khalique's avatar
Khalique committed
569
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
570

Khalique's avatar
Khalique committed
571

Khalique's avatar
Khalique committed
572
@onnx_test
Khalique's avatar
Khalique committed
573
574
575
576
577
578
579
580
581
582
def erf_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 15])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10, 15])

    node = onnx.helper.make_node(
        'Erf',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
583
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
584

Khalique's avatar
Khalique committed
585

Khalique's avatar
Khalique committed
586
@onnx_test
Khalique's avatar
Khalique committed
587
588
589
590
591
592
593
594
595
596
def exp_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Exp',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
597
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
598

Khalique's avatar
Khalique committed
599

Khalique's avatar
Khalique committed
600
@onnx_test
Khalique's avatar
Khalique committed
601
602
def expand_test():
    shape_val = np.array([2, 3, 4, 5]).astype(np.int64)
Khalique's avatar
Khalique committed
603
604
605
606
    shape_ts = helper.make_tensor(name='shape_tensor',
                                  data_type=TensorProto.INT32,
                                  dims=shape_val.shape,
                                  vals=shape_val.flatten().astype(int))
Khalique's avatar
Khalique committed
607
608
609
610
611
612
613
614
615
    shape_const = helper.make_node(
        'Constant',
        inputs=[],
        outputs=['shape'],
        value=shape_ts,
    )
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 1, 1])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 3, 4, 5])

Khalique's avatar
Khalique committed
616
617
618
619
620
621
    node = onnx.helper.make_node('Expand',
                                 inputs=['x', 'shape'],
                                 outputs=['y'])

    return ([shape_const, node], [x], [y])

Khalique's avatar
Khalique committed
622

Khalique's avatar
Khalique committed
623
@onnx_test
Khalique's avatar
Khalique committed
624
625
def flatten_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
Khalique's avatar
Khalique committed
626
    y = helper.make_tensor_value_info('2', TensorProto.FLOAT, [6, 20])
Khalique's avatar
Khalique committed
627
628
    y2 = helper.make_tensor_value_info('3', TensorProto.FLOAT, [2, 60])

Khalique's avatar
Khalique committed
629
630
631
632
    node = onnx.helper.make_node('Flatten',
                                 inputs=['0'],
                                 axis=2,
                                 outputs=['2'])
Khalique's avatar
Khalique committed
633

Khalique's avatar
Khalique committed
634
635
636
    node2 = onnx.helper.make_node('Flatten', inputs=['0'], outputs=['3'])

    return ([node, node2], [x], [y, y2])
Khalique's avatar
Khalique committed
637
638


Khalique's avatar
Khalique committed
639
@onnx_test
Khalique's avatar
Khalique committed
640
641
def gather_test():
    x = helper.make_tensor_value_info('data', TensorProto.FLOAT, [3, 4, 5, 6])
Khalique's avatar
Khalique committed
642
643
    i = helper.make_tensor_value_info('indices', TensorProto.INT32,
                                      [2, 3, 4, 5])
Khalique's avatar
Khalique committed
644
645
646
647
648
649
650
651
652
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 3, 4, 5])

    node = onnx.helper.make_node(
        'Gather',
        inputs=['data', 'indices'],
        outputs=['y'],
        axis=1,
    )

Khalique's avatar
Khalique committed
653
654
    return ([node], [x, i], [y])

Khalique's avatar
Khalique committed
655

Khalique's avatar
Khalique committed
656
@onnx_test
Khalique's avatar
Khalique committed
657
658
659
660
661
662
def gemm_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [5, 7])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [11, 5])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [])
    a = helper.make_tensor_value_info('3', TensorProto.FLOAT, [7, 11])

Khalique's avatar
Khalique committed
663
664
665
666
667
668
669
670
671
    node = onnx.helper.make_node('Gemm',
                                 inputs=['0', '1', '2'],
                                 outputs=['3'],
                                 alpha=2.0,
                                 beta=2.0,
                                 transA=1,
                                 transB=1)

    return ([node], [x, y, z], [a])
Khalique's avatar
Khalique committed
672
673


Khalique's avatar
Khalique committed
674
@onnx_test
Khalique's avatar
Khalique committed
675
676
677
678
679
680
def gemm_ex_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 1, 5, 6])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1, 1, 5, 7])
    m3 = helper.make_tensor_value_info('3', TensorProto.FLOAT, [1, 1, 6, 7])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 1, 6, 7])

Khalique's avatar
Khalique committed
681
682
683
684
685
686
687
688
    node = onnx.helper.make_node('Gemm',
                                 inputs=['1', '2', '3'],
                                 outputs=['y'],
                                 alpha=0.5,
                                 beta=0.8,
                                 transA=1)

    return ([node], [m1, m2, m3], [y])
Khalique's avatar
Khalique committed
689
690


Khalique's avatar
Khalique committed
691
@onnx_test
Khalique's avatar
Khalique committed
692
693
694
695
696
697
def gemm_ex_brcst_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 1, 5, 6])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1, 1, 5, 7])
    m3 = helper.make_tensor_value_info('3', TensorProto.FLOAT, [1, 1, 6, 1])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 1, 6, 7])

Khalique's avatar
Khalique committed
698
699
700
701
702
703
704
705
    node = onnx.helper.make_node('Gemm',
                                 inputs=['1', '2', '3'],
                                 outputs=['y'],
                                 alpha=0.5,
                                 beta=0.8,
                                 transA=1)

    return ([node], [m1, m2, m3], [y])
Khalique's avatar
Khalique committed
706
707


Khalique's avatar
Khalique committed
708
@onnx_test
Khalique's avatar
Khalique committed
709
def globalavgpool_test():
Khalique's avatar
Khalique committed
710
711
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 16, 16])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 1, 1])
Khalique's avatar
Khalique committed
712
713
714
715
716
717
718

    node = onnx.helper.make_node(
        'GlobalAveragePool',
        inputs=['0'],
        outputs=['1'],
    )

Khalique's avatar
Khalique committed
719
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
720

Khalique's avatar
Khalique committed
721

Khalique's avatar
Khalique committed
722
@onnx_test
Khalique's avatar
Khalique committed
723
def globalmaxpool_test():
Khalique's avatar
Khalique committed
724
725
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 16, 16])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 1, 1])
Khalique's avatar
Khalique committed
726
727
728
729
730
731
732

    node = onnx.helper.make_node(
        'GlobalMaxPool',
        inputs=['0'],
        outputs=['1'],
    )

Khalique's avatar
Khalique committed
733
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
734

Khalique's avatar
Khalique committed
735

Khalique's avatar
Khalique committed
736
@onnx_test
Khalique's avatar
Khalique committed
737
738
739
740
741
742
743
744
745
746
747
748
def group_conv_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 4, 16, 16])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [4, 1, 3, 3])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [1, 4, 14, 14])

    node = onnx.helper.make_node(
        'Conv',
        inputs=['0', '1'],
        group=4,
        outputs=['2'],
    )

Khalique's avatar
Khalique committed
749
750
    return ([node], [x, y], [z])

Khalique's avatar
Khalique committed
751

Khalique's avatar
Khalique committed
752
@onnx_test
Khalique's avatar
Khalique committed
753
def imagescaler_test():
Khalique's avatar
Khalique committed
754
755
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3, 16, 16])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 16, 16])
Khalique's avatar
Khalique committed
756

Khalique's avatar
Khalique committed
757
758
759
760
761
    node = onnx.helper.make_node('ImageScaler',
                                 inputs=['0'],
                                 outputs=['1'],
                                 bias=[0.01, 0.02, 0.03],
                                 scale=0.5)
Khalique's avatar
Khalique committed
762

Khalique's avatar
Khalique committed
763
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
764

Khalique's avatar
Khalique committed
765

Khalique's avatar
Khalique committed
766
@onnx_test
Khalique's avatar
Khalique committed
767
768
769
770
771
772
773
774
775
776
777
def implicit_add_bcast_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 4, 1])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [2, 3, 4, 5])

    node = onnx.helper.make_node(
        'Add',
        inputs=['0', '1'],
        outputs=['2'],
    )

Khalique's avatar
Khalique committed
778
779
    return ([node], [x, y], [z])

Khalique's avatar
Khalique committed
780

Khalique's avatar
Khalique committed
781
@onnx_test
Khalique's avatar
Khalique committed
782
783
784
def implicit_pow_bcast_test():
    arg0 = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    arg1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 4, 1])
Khalique's avatar
Khalique committed
785
786
    arg_out = helper.make_tensor_value_info('out', TensorProto.FLOAT,
                                            [2, 3, 4, 5])
Khalique's avatar
Khalique committed
787
788
789
790
791
792
793

    node = onnx.helper.make_node(
        'Pow',
        inputs=['0', '1'],
        outputs=['out'],
    )

Khalique's avatar
Khalique committed
794
795
    return ([node], [arg0, arg1], [arg_out])

Khalique's avatar
Khalique committed
796

Khalique's avatar
Khalique committed
797
@onnx_test
Khalique's avatar
Khalique committed
798
799
800
def implicit_sub_bcast_test():
    arg0 = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    arg1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [4, 5])
Khalique's avatar
Khalique committed
801
802
    arg_out = helper.make_tensor_value_info('out', TensorProto.FLOAT,
                                            [2, 3, 4, 5])
Khalique's avatar
Khalique committed
803
804
805
806
807
808
809

    node = onnx.helper.make_node(
        'Sub',
        inputs=['0', '1'],
        outputs=['out'],
    )

Khalique's avatar
Khalique committed
810
811
    return ([node], [arg0, arg1], [arg_out])

Khalique's avatar
Khalique committed
812

Khalique's avatar
Khalique committed
813
@onnx_test
Khalique's avatar
Khalique committed
814
815
816
817
def leaky_relu_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])

Khalique's avatar
Khalique committed
818
819
820
821
    node = onnx.helper.make_node('LeakyRelu',
                                 inputs=['0'],
                                 outputs=['1'],
                                 alpha=0.01)
Khalique's avatar
Khalique committed
822

Khalique's avatar
Khalique committed
823
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
824

Khalique's avatar
Khalique committed
825

Khalique's avatar
Khalique committed
826
@onnx_test
Khalique's avatar
Khalique committed
827
828
829
830
831
832
833
834
835
836
def log_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Log',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
837
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
838

Khalique's avatar
Khalique committed
839

Khalique's avatar
Khalique committed
840
@onnx_test
Khalique's avatar
Khalique committed
841
842
843
844
def logsoftmax_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 5, 6])

Khalique's avatar
Khalique committed
845
846
847
848
    node = onnx.helper.make_node('LogSoftmax',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axis=1)
Khalique's avatar
Khalique committed
849

Khalique's avatar
Khalique committed
850
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
851

Khalique's avatar
Khalique committed
852

Khalique's avatar
Khalique committed
853
@onnx_test
Khalique's avatar
Khalique committed
854
855
856
857
def lrn_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 28, 24, 24])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 28, 24, 24])

Khalique's avatar
Khalique committed
858
859
860
861
862
863
864
    node = onnx.helper.make_node('LRN',
                                 inputs=['0'],
                                 size=5,
                                 alpha=0.0001,
                                 beta=0.75,
                                 bias=1.0,
                                 outputs=['1'])
Khalique's avatar
Khalique committed
865

Khalique's avatar
Khalique committed
866
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
867

Khalique's avatar
Khalique committed
868

Khalique's avatar
Khalique committed
869
@onnx_test
Khalique's avatar
Khalique committed
870
871
872
873
874
875
876
877
878
879
880
def matmul_bmbm_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 6, 7])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [5, 2, 1, 7, 8])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [5, 2, 3, 6, 8])

    node = onnx.helper.make_node(
        'MatMul',
        inputs=['1', '2'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
881
882
    return ([node], [m1, m2], [y])

Khalique's avatar
Khalique committed
883

Khalique's avatar
Khalique committed
884
@onnx_test
Khalique's avatar
Khalique committed
885
886
887
888
889
890
891
892
893
894
895
def matmul_bmv_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 6, 7])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [7])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 6])

    node = onnx.helper.make_node(
        'MatMul',
        inputs=['1', '2'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
896
897
    return ([node], [m1, m2], [y])

Khalique's avatar
Khalique committed
898

Khalique's avatar
Khalique committed
899
@onnx_test
Khalique's avatar
Khalique committed
900
901
902
903
904
905
906
907
908
909
910
def matmul_mv_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [6, 7])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [7])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [6])

    node = onnx.helper.make_node(
        'MatMul',
        inputs=['1', '2'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
911
912
    return ([node], [m1, m2], [y])

Khalique's avatar
Khalique committed
913

Khalique's avatar
Khalique committed
914
@onnx_test
Khalique's avatar
Khalique committed
915
916
917
918
919
920
921
922
923
924
925
def matmul_vbm_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [7])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [5, 7, 8])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [5, 8])

    node = onnx.helper.make_node(
        'MatMul',
        inputs=['1', '2'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
926
927
    return ([node], [m1, m2], [y])

Khalique's avatar
Khalique committed
928

Khalique's avatar
Khalique committed
929
@onnx_test
Khalique's avatar
Khalique committed
930
931
932
933
934
935
936
937
938
939
940
def matmul_vm_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [7])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [7, 8])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [8])

    node = onnx.helper.make_node(
        'MatMul',
        inputs=['1', '2'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
941
942
    return ([node], [m1, m2], [y])

Khalique's avatar
Khalique committed
943

Khalique's avatar
Khalique committed
944
@onnx_test
Khalique's avatar
Khalique committed
945
946
947
948
949
950
951
952
953
954
955
def matmul_vv_test():
    m1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [7])
    m2 = helper.make_tensor_value_info('2', TensorProto.FLOAT, [7])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1])

    node = onnx.helper.make_node(
        'MatMul',
        inputs=['1', '2'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
956
957
    return ([node], [m1, m2], [y])

Khalique's avatar
Khalique committed
958

Khalique's avatar
Khalique committed
959
@onnx_test
Khalique's avatar
Khalique committed
960
961
962
963
964
965
966
967
968
969
970
971
def max_test():
    a = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    b = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])
    c = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3])
    y = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3])

    node = onnx.helper.make_node(
        'Max',
        inputs=['0', '1', '2'],
        outputs=['3'],
    )

Khalique's avatar
Khalique committed
972
973
    return ([node], [a, b, c], [y])

Khalique's avatar
Khalique committed
974

Khalique's avatar
Khalique committed
975
@onnx_test
Khalique's avatar
Khalique committed
976
977
978
979
980
981
982
983
984
985
986
987
def min_test():
    a = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    b = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])
    c = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3])
    y = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3])

    node = onnx.helper.make_node(
        'Min',
        inputs=['0', '1', '2'],
        outputs=['3'],
    )

Khalique's avatar
Khalique committed
988
989
    return ([node], [a, b, c], [y])

Khalique's avatar
Khalique committed
990

Khalique's avatar
Khalique committed
991
@onnx_test
Khalique's avatar
Khalique committed
992
993
994
995
def no_pad_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 2])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [2, 2])

Khalique's avatar
Khalique committed
996
997
998
999
    node = onnx.helper.make_node('Pad',
                                 inputs=['0'],
                                 pads=[0, 0, 0, 0],
                                 outputs=['1'])
Khalique's avatar
Khalique committed
1000

Khalique's avatar
Khalique committed
1001
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1002

Khalique's avatar
Khalique committed
1003

Khalique's avatar
Khalique committed
1004
@onnx_test
Khalique's avatar
Khalique committed
1005
1006
1007
1008
def pad_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 2])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [4, 4])

Khalique's avatar
Khalique committed
1009
1010
1011
1012
    node = onnx.helper.make_node('Pad',
                                 inputs=['0'],
                                 pads=[1, 1, 1, 1],
                                 outputs=['1'])
Khalique's avatar
Khalique committed
1013

Khalique's avatar
Khalique committed
1014
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1015

Khalique's avatar
Khalique committed
1016

Khalique's avatar
Khalique committed
1017
@onnx_test
Khalique's avatar
Khalique committed
1018
1019
1020
def pow_test():
    arg0 = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    arg1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [2, 3, 4, 5])
Khalique's avatar
Khalique committed
1021
1022
    arg_out = helper.make_tensor_value_info('out', TensorProto.FLOAT,
                                            [2, 3, 4, 5])
Khalique's avatar
Khalique committed
1023
1024
1025
1026
1027
1028
1029

    node = onnx.helper.make_node(
        'Pow',
        inputs=['0', '1'],
        outputs=['out'],
    )

Khalique's avatar
Khalique committed
1030
    return ([node], [arg0, arg1], [arg_out])
Khalique's avatar
Khalique committed
1031

Khalique's avatar
Khalique committed
1032

Khalique's avatar
Khalique committed
1033
@onnx_test
Khalique's avatar
Khalique committed
1034
1035
1036
def reducemean_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4])
Khalique's avatar
Khalique committed
1037
    axes = [2, 3]
Khalique's avatar
Khalique committed
1038

Khalique's avatar
Khalique committed
1039
1040
1041
1042
1043
    node = onnx.helper.make_node('ReduceMean',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axes=axes,
                                 keepdims=0)
Khalique's avatar
Khalique committed
1044

Khalique's avatar
Khalique committed
1045
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1046

Khalique's avatar
Khalique committed
1047

Khalique's avatar
Khalique committed
1048
@onnx_test
Khalique's avatar
Khalique committed
1049
1050
1051
def reducemean_keepdims_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 1, 6])
Khalique's avatar
Khalique committed
1052
    axes = [2]
Khalique's avatar
Khalique committed
1053

Khalique's avatar
Khalique committed
1054
1055
1056
1057
1058
    node = onnx.helper.make_node('ReduceMean',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axes=axes,
                                 keepdims=1)
Khalique's avatar
Khalique committed
1059

Khalique's avatar
Khalique committed
1060
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1061

Khalique's avatar
Khalique committed
1062

Khalique's avatar
Khalique committed
1063
@onnx_test
Khalique's avatar
Khalique committed
1064
def reducesum_test():
Khalique's avatar
Khalique committed
1065
1066
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 1, 1])
Khalique's avatar
Khalique committed
1067
    axes = [2]
Khalique's avatar
Khalique committed
1068

Khalique's avatar
Khalique committed
1069
1070
1071
1072
1073
    node = onnx.helper.make_node('ReduceSum',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axes=axes,
                                 keepdims=0)
Khalique's avatar
Khalique committed
1074

Khalique's avatar
Khalique committed
1075
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1076

Khalique's avatar
Khalique committed
1077

Khalique's avatar
Khalique committed
1078
@onnx_test
Khalique's avatar
Khalique committed
1079
1080
1081
def reducesum_multiaxis_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 1, 1])
Khalique's avatar
Khalique committed
1082
    axes = [2, 3]
Khalique's avatar
Khalique committed
1083

Khalique's avatar
Khalique committed
1084
1085
1086
1087
1088
    node = onnx.helper.make_node('ReduceSum',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axes=axes,
                                 keepdims=0)
Khalique's avatar
Khalique committed
1089

Khalique's avatar
Khalique committed
1090
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1091

Khalique's avatar
Khalique committed
1092

Khalique's avatar
Khalique committed
1093
@onnx_test
Khalique's avatar
Khalique committed
1094
1095
1096
def reducesum_keepdims_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4, 1, 1])
Khalique's avatar
Khalique committed
1097
    axes = [2, 3]
Khalique's avatar
Khalique committed
1098

Khalique's avatar
Khalique committed
1099
1100
1101
1102
1103
    node = onnx.helper.make_node('ReduceSum',
                                 inputs=['x'],
                                 outputs=['y'],
                                 axes=axes,
                                 keepdims=1)
Khalique's avatar
Khalique committed
1104

Khalique's avatar
Khalique committed
1105
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1106

Khalique's avatar
Khalique committed
1107

Khalique's avatar
Khalique committed
1108
@onnx_test
Khalique's avatar
Khalique committed
1109
1110
1111
def reshape_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [4, 2, 3])
    x_shape = helper.make_tensor_value_info('1', TensorProto.INT64, [2])
Khalique's avatar
Khalique committed
1112
    x_shape_list = [3, 8]
Khalique's avatar
Khalique committed
1113
1114
1115
    y = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3, 8])
    y2 = helper.make_tensor_value_info('3', TensorProto.FLOAT, [3, 8])

Khalique's avatar
Khalique committed
1116
    node = onnx.helper.make_node('Reshape', inputs=['0', '1'], outputs=['2'])
Khalique's avatar
Khalique committed
1117

Khalique's avatar
Khalique committed
1118
1119
1120
1121
1122
1123
1124
    node2 = onnx.helper.make_node('Reshape',
                                  inputs=['0'],
                                  shape=x_shape_list,
                                  outputs=['3'])

    return ([node, node2], [x, x_shape], [y, y2],
            [helper.make_tensor('1', TensorProto.INT64, [2], [3, 8])])
Khalique's avatar
Khalique committed
1125
1126


Khalique's avatar
Khalique committed
1127
@onnx_test
Khalique's avatar
Khalique committed
1128
1129
def reshape_non_standard_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [2, 3, 4])
Khalique's avatar
Khalique committed
1130
1131
    trans_x = helper.make_tensor_value_info('trans_x', TensorProto.FLOAT,
                                            [2, 4, 3])
Khalique's avatar
Khalique committed
1132
1133
1134
1135
1136
1137
1138
1139
1140
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [4, 3, 2])

    trans = helper.make_node(
        'Transpose',
        inputs=['x'],
        outputs=['trans_x'],
        perm=[0, 2, 1],
    )

Khalique's avatar
Khalique committed
1141
1142
1143
1144
1145
1146
    res = onnx.helper.make_node('Reshape',
                                inputs=['trans_x'],
                                outputs=['y'],
                                shape=[4, 3, 2])

    return ([trans, res], [x], [y])
Khalique's avatar
Khalique committed
1147
1148


Khalique's avatar
Khalique committed
1149
@onnx_test
Khalique's avatar
Khalique committed
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
def shape_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4, 5, 6])
    y = helper.make_tensor_value_info('y', TensorProto.INT64, [4])

    node = onnx.helper.make_node(
        'Shape',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
1160
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1161

Khalique's avatar
Khalique committed
1162

Khalique's avatar
Khalique committed
1163
@onnx_test
Khalique's avatar
Khalique committed
1164
1165
1166
1167
1168
1169
1170
def shape_gather_test():
    values = np.array([1])
    value = helper.make_tensor_value_info('value', TensorProto.INT32, [1])
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [7, 3, 10])
    y = helper.make_tensor_value_info('y', TensorProto.INT64, [3])
    z = helper.make_tensor_value_info('z', TensorProto.FLOAT, [1])

Khalique's avatar
Khalique committed
1171
1172
1173
1174
    value_tensor = helper.make_tensor(name='const_tensor',
                                      data_type=TensorProto.INT32,
                                      dims=values.shape,
                                      vals=values.flatten().astype(int))
Khalique's avatar
Khalique committed
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195

    node_const = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['value'],
        value=value_tensor,
    )

    node_shape = onnx.helper.make_node(
        'Shape',
        inputs=['x'],
        outputs=['y'],
    )

    node_gather = helper.make_node(
        'Gather',
        inputs=['y', 'value'],
        outputs=['z'],
        axis=0,
    )

Khalique's avatar
Khalique committed
1196
1197
    return ([node_const, node_shape, node_gather], [x], [z])

Khalique's avatar
Khalique committed
1198

Khalique's avatar
Khalique committed
1199
@onnx_test
Khalique's avatar
Khalique committed
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
def sign_test():
    x = helper.make_tensor_value_info('x', TensorProto.DOUBLE, [10, 5])
    y = helper.make_tensor_value_info('y', TensorProto.DOUBLE, [10, 5])

    node = onnx.helper.make_node(
        'Sign',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
1210
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1211

Khalique's avatar
Khalique committed
1212

Khalique's avatar
Khalique committed
1213
@onnx_test
Khalique's avatar
Khalique committed
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
def sin_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Sin',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
1224
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1225

Khalique's avatar
Khalique committed
1226

Khalique's avatar
Khalique committed
1227
@onnx_test
Khalique's avatar
Khalique committed
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
def sinh_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
        'Sinh',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
1238
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1239

Khalique's avatar
Khalique committed
1240

Khalique's avatar
Khalique committed
1241
@onnx_test
Khalique's avatar
Khalique committed
1242
1243
1244
1245
def slice_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3, 2])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 2])

Khalique's avatar
Khalique committed
1246
1247
1248
1249
1250
1251
    node = onnx.helper.make_node('Slice',
                                 inputs=['0'],
                                 axes=[0, 1],
                                 starts=[1, 0],
                                 ends=[2, 2],
                                 outputs=['1'])
Khalique's avatar
Khalique committed
1252

Khalique's avatar
Khalique committed
1253
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1254

Khalique's avatar
Khalique committed
1255

Khalique's avatar
Khalique committed
1256
@onnx_test
Khalique's avatar
Khalique committed
1257
1258
1259
1260
def softmax_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 3])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3])

Khalique's avatar
Khalique committed
1261
    node = onnx.helper.make_node('Softmax', inputs=['0'], outputs=['1'])
Khalique's avatar
Khalique committed
1262

Khalique's avatar
Khalique committed
1263
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1264

Khalique's avatar
Khalique committed
1265

Khalique's avatar
Khalique committed
1266
@onnx_test
Khalique's avatar
Khalique committed
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
def sqrt_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 15])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10, 15])

    node = onnx.helper.make_node(
        'Sqrt',
        inputs=['x'],
        outputs=['y'],
    )

Khalique's avatar
Khalique committed
1277
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1278

Khalique's avatar
Khalique committed
1279

Khalique's avatar
Khalique committed
1280
@onnx_test
Khalique's avatar
Khalique committed
1281
def squeeze_unsqueeze_test():
Khalique's avatar
Khalique committed
1282
1283
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT,
                                      [1, 3, 1, 1, 2, 1])
Khalique's avatar
Khalique committed
1284
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 2])
Khalique's avatar
Khalique committed
1285
1286
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT,
                                      [1, 1, 3, 1, 2, 1])
Khalique's avatar
Khalique committed
1287

Khalique's avatar
Khalique committed
1288
1289
1290
1291
    node = onnx.helper.make_node('Squeeze',
                                 inputs=['0'],
                                 axes=[0, 2, 3, 5],
                                 outputs=['1'])
Khalique's avatar
Khalique committed
1292

Khalique's avatar
Khalique committed
1293
1294
1295
1296
1297
1298
    node2 = onnx.helper.make_node('Unsqueeze',
                                  inputs=['1'],
                                  axes=[0, 1, 3, 5],
                                  outputs=['2'])

    return ([node, node2], [x], [z])
Khalique's avatar
Khalique committed
1299
1300


Khalique's avatar
Khalique committed
1301
@onnx_test
Khalique's avatar
Khalique committed
1302
1303
1304
def sub_bcast_test():
    arg0 = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    arg1 = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 4])
Khalique's avatar
Khalique committed
1305
1306
    arg_out = helper.make_tensor_value_info('out', TensorProto.FLOAT,
                                            [2, 3, 4, 5])
Khalique's avatar
Khalique committed
1307
1308
1309
1310
1311

    node = onnx.helper.make_node(
        'Sub',
        inputs=['0', '1'],
        outputs=['out'],
Khalique's avatar
Khalique committed
1312
1313
        broadcast=1,
        axis=1,
Khalique's avatar
Khalique committed
1314
1315
    )

Khalique's avatar
Khalique committed
1316
1317
    return ([node], [arg0, arg1], [arg_out])

Khalique's avatar
Khalique committed
1318

Khalique's avatar
Khalique committed
1319
@onnx_test
Khalique's avatar
Khalique committed
1320
1321
def sub_scalar_test():
    values = np.array([1])
Khalique's avatar
Khalique committed
1322
1323
1324
1325
1326
1327
1328
1329
1330
    arg_node = helper.make_tensor_value_info('0', TensorProto.FLOAT,
                                             [2, 3, 4, 5])
    arg_out = helper.make_tensor_value_info('out', TensorProto.FLOAT,
                                            [2, 3, 4, 5])

    values_tensor = helper.make_tensor(name='const',
                                       data_type=TensorProto.FLOAT,
                                       dims=values.shape,
                                       vals=values.flatten().astype(float))
Khalique's avatar
Khalique committed
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344

    arg_const = onnx.helper.make_node(
        'Constant',
        inputs=[],
        outputs=['arg_const'],
        value=values_tensor,
    )

    node = onnx.helper.make_node(
        'Sub',
        inputs=['0', 'arg_const'],
        outputs=['out'],
    )

Khalique's avatar
Khalique committed
1345
1346
    return ([arg_const, node], [arg_node], [arg_out])

Khalique's avatar
Khalique committed
1347

Khalique's avatar
Khalique committed
1348
@onnx_test
Khalique's avatar
Khalique committed
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
def sum_test():
    a = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    b = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])
    c = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3])

    y = helper.make_tensor_value_info('3', TensorProto.FLOAT, [3])

    node = onnx.helper.make_node(
        'Sum',
        inputs=['0', '1', '2'],
        outputs=['3'],
    )

Khalique's avatar
Khalique committed
1362
1363
    return ([node], [a, b, c], [y])

Khalique's avatar
Khalique committed
1364

Khalique's avatar
Khalique committed
1365
@onnx_test
Khalique's avatar
Khalique committed
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
def sum_test():
    a = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
    b = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3])
    c = helper.make_tensor_value_info('2', TensorProto.FLOAT, [3])
    y = helper.make_tensor_value_info('3', TensorProto.FLOAT, [3])

    node = onnx.helper.make_node(
        'Sum',
        inputs=['0', '1', '2'],
        outputs=['3'],
    )

Khalique's avatar
Khalique committed
1378
1379
    return ([node], [a, b, c], [y])

Khalique's avatar
Khalique committed
1380

Khalique's avatar
Khalique committed
1381
@onnx_test
Khalique's avatar
Khalique committed
1382
1383
1384
1385
1386
def tan_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [10])

    node = onnx.helper.make_node(
Khalique's avatar
Khalique committed
1387
1388
1389
1390
        'Tan',
        inputs=['x'],
        outputs=['y'],
    )
Khalique's avatar
Khalique committed
1391

Khalique's avatar
Khalique committed
1392
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1393

Khalique's avatar
Khalique committed
1394

Khalique's avatar
Khalique committed
1395
@onnx_test
Khalique's avatar
Khalique committed
1396
1397
1398
1399
1400
def tanh_test():
    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1])

    node = onnx.helper.make_node(
Khalique's avatar
Khalique committed
1401
1402
1403
1404
        'Tanh',
        inputs=['x'],
        outputs=['y'],
    )
Khalique's avatar
Khalique committed
1405

Khalique's avatar
Khalique committed
1406
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1407

Khalique's avatar
Khalique committed
1408

Khalique's avatar
Khalique committed
1409
@onnx_test
Khalique's avatar
Khalique committed
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
def transpose_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 2, 2, 3])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [1, 3, 2, 2])

    node = onnx.helper.make_node(
        'Transpose',
        perm=[0, 3, 1, 2],
        inputs=['0'],
        outputs=['1'],
    )

Khalique's avatar
Khalique committed
1421
    return ([node], [x], [y])
Khalique's avatar
Khalique committed
1422

Khalique's avatar
Khalique committed
1423

Khalique's avatar
Khalique committed
1424
1425
1426
@onnx_test
def transpose_gather_test():
    x = helper.make_tensor_value_info('data', TensorProto.FLOAT, [3, 5, 4, 6])
Khalique's avatar
Khalique committed
1427
1428
1429
1430
    i = helper.make_tensor_value_info('indices', TensorProto.INT32,
                                      [2, 4, 3, 5])
    y = helper.make_tensor_value_info('y', TensorProto.FLOAT,
                                      [3, 2, 3, 4, 5, 4, 5, 6])
Khalique's avatar
Khalique committed
1431
1432
1433
1434
1435
1436
1437
1438

    td = onnx.helper.make_node(
        'Transpose',
        inputs=['data'],
        outputs=['tdata'],
        perm=[0, 2, 1, 3],
    )

Khalique's avatar
Khalique committed
1439
1440
1441
1442
    ti = onnx.helper.make_node('Transpose',
                               inputs=['indices'],
                               outputs=['tindices'],
                               perm=[0, 2, 1, 3])
Khalique's avatar
Khalique committed
1443
1444
1445
1446
1447
1448
1449
1450

    node = onnx.helper.make_node(
        'Gather',
        inputs=['tdata', 'tindices'],
        outputs=['y'],
        axis=1,
    )

Khalique's avatar
Khalique committed
1451
    return ([td, ti, node], [x, i], [y])
Khalique's avatar
Khalique committed
1452

Khalique's avatar
Khalique committed
1453

Khalique's avatar
Khalique committed
1454
@onnx_test
Khalique's avatar
Khalique committed
1455
1456
1457
1458
1459
1460
def unknown_test():
    x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])
    y = helper.make_tensor_value_info('1', TensorProto.FLOAT, [3, 4])
    z = helper.make_tensor_value_info('2', TensorProto.FLOAT, [2, 3, 4, 5])
    a = helper.make_tensor_value_info('3', TensorProto.FLOAT, [2, 3, 4, 5])

Khalique's avatar
Khalique committed
1461
    node = onnx.helper.make_node('Unknown', inputs=['0', '1'], outputs=['2'])
Khalique's avatar
Khalique committed
1462

Khalique's avatar
Khalique committed
1463
    node2 = onnx.helper.make_node('Unknown', inputs=['2'], outputs=['3'])
Khalique's avatar
Khalique committed
1464

Khalique's avatar
Khalique committed
1465
    return ([node, node2], [x, y], [a])