onnx_conditional_example.py

import onnx
import numpy as np


from onnx import numpy_helper, TensorProto, save_model
from onnx.helper import make_model, make_node, make_graph, make_tensor_value_info
from onnx.checker import check_model

# condition has to be a float tensor
condition = make_tensor_value_info('condition', TensorProto.FLOAT, [1])
input1 = make_tensor_value_info('input1', TensorProto.FLOAT, [1])
input2 = make_tensor_value_info('input2', TensorProto.FLOAT, [1])
output = make_tensor_value_info('output', TensorProto.FLOAT, [1])

then_branch = make_graph(
    nodes=[
        make_node('Add', ['input1', 'input2'], ['output'])
    ],
    name='then_branch',
    inputs=[input1, input2],
    outputs=[output]
)

else_branch = make_graph(
    nodes=[
        make_node('Sub', ['input1', 'input2'], ['output'])
    ],
    name='else_branch',
    inputs=[input1, input2],
    outputs=[output]
)

graph = make_graph(
    nodes=[
        make_node('If', ['condition'], ['output'], then_branch=then_branch, else_branch=else_branch)
    ],
    name='if_example',
    inputs=[condition],
    outputs=[output]
)

model = make_model(graph, producer_name='conditional_example')

# Check the model and save it
check_model(model)

save_model(model, 'conditional_example.onnx')

# Load the model and check it
model = onnx.load('conditional_example.onnx')
check_model(model)