diff options
Diffstat (limited to 'searchlib/src/tests/features/onnx_feature/fragile.py')
-rwxr-xr-x | searchlib/src/tests/features/onnx_feature/fragile.py | 30 |
1 files changed, 30 insertions, 0 deletions
diff --git a/searchlib/src/tests/features/onnx_feature/fragile.py b/searchlib/src/tests/features/onnx_feature/fragile.py new file mode 100755 index 00000000000..e4eaf168e14 --- /dev/null +++ b/searchlib/src/tests/features/onnx_feature/fragile.py @@ -0,0 +1,30 @@ +# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +import onnx +from onnx import helper, TensorProto + +INPUT1 = helper.make_tensor_value_info('in1', TensorProto.FLOAT, [2]) +INPUT2 = helper.make_tensor_value_info('in2', TensorProto.FLOAT, ['batch']) + +OUTPUT = helper.make_tensor_value_info('out', TensorProto.FLOAT, ['batch']) + +nodes = [ + helper.make_node( + 'Add', + ['in1', 'in2'], + ['out'], + ), +] +graph_def = helper.make_graph( + nodes, + 'fragile', + [ + INPUT1, + INPUT2, + ], + [ + OUTPUT, + ], +) +model_def = helper.make_model(graph_def, producer_name='fragile.py', opset_imports=[onnx.OperatorSetIdProto(version=12)]) +onnx.save(model_def, 'fragile.onnx') |