onnx_loop_with_scan_outputs_example.py

import onnx
from onnx import helper, TensorProto, save_model
from onnx.checker import check_model

body = helper.make_graph(
    name='body',
    nodes=[
        helper.make_node('Add', ['a', 'b_in'], ['my_local']),
        helper.make_node('Sub', ['a', 'b_in'], ['b_out']),
        helper.make_node('Greater', ['my_local', 'b_out'], ['keep_going_out']),
        helper.make_node('Add', ['b_in', 'b_in'], ['user_defined_vals']),
    ],
    inputs=[
        helper.make_tensor_value_info('iteration_num', TensorProto.INT64, [1]),
        helper.make_tensor_value_info('keep_going_in', TensorProto.BOOL, [1]),
        helper.make_tensor_value_info('b_in', TensorProto.FLOAT, [1])
    ],
    outputs=[
        helper.make_tensor_value_info('keep_going_out', TensorProto.BOOL, [1]),
        helper.make_tensor_value_info('b_out', TensorProto.FLOAT, [1]),
        helper.make_tensor_value_info('user_defined_vals', TensorProto.FLOAT, [1]),
    ]
)

graph = helper.make_graph(
    name='graph',
    nodes=[
        helper.make_node('Cast', ['keep_going'], ['keep_going_bool'], to=TensorProto.BOOL),
        helper.make_node('Loop', ['max_trip_count', 'keep_going_bool', 'b'],
                         ['b_loop', 'user_defined_vals_loop'], body=body),
    ],
    inputs=[],
    outputs=[
        helper.make_tensor_value_info('b_loop', TensorProto.FLOAT, [1]),
        helper.make_tensor_value_info('user_defined_vals_loop', TensorProto.FLOAT, [-1, 1]),
    ],
    initializer=[
        helper.make_tensor('a', TensorProto.FLOAT, [1], [3]),
        helper.make_tensor('b', TensorProto.FLOAT, [1], [6]),
        helper.make_tensor('keep_going', TensorProto.UINT8, [1], [True]),
        helper.make_tensor('max_trip_count', TensorProto.INT64, [1], [10]),
    ],
)

model = helper.make_model(graph, opset_imports=[onnx.OperatorSetIdProto(version=21)])
check_model(model)
save_model(model, 'loop_with_scan_outputs_example.onnx')