unsorted_subgraph.py

import onnx_ir as ir
import onnx


def create_model():
    """Create a model that has a unsorted node with subgraph that uses a value defined later."""
    a = ir.Value(name="a", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((2, 3)))
    b = ir.Value(name="b", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((3, 4)))
    b_out = ir.Value(name="b_out", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((3, 4)))
    c = ir.Value(name="c", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape((4, 5)))

    b_producer = ir.node(
        name="b_producer",
        op_type="SomeOp",
        inputs=[a],
        outputs=[b],
    )

    node_with_subgraph = ir.node(
        name="node_with_subgraph",
        op_type="SubgraphOp",
        inputs=[],
        outputs=[c],
        attributes={
            "subgraph": ir.Graph(
                inputs=[],
                outputs=[],
                nodes=[
                    ir.node(
                        name="subgraph_node",
                        op_type="SomeOp",
                        inputs=[b],
                        outputs=[b_out],
                    )
                ],
                name="subgraph",
            )
        }
    )

    main_graph = ir.Graph(
        inputs=[a],
        outputs=[c],
        nodes=[node_with_subgraph, b_producer],
        opset_imports={
            "": 42
        },
        name="main_graph",
    )
    model = ir.Model(main_graph, ir_version=10)

    model_text = ir.to_onnx_text(model)
    print(model_text)

    onnx.checker.check_model(ir.to_proto(model))

    m = ir.from_onnx_text(
        model_text
    )
    m.graph.sort()
    print(ir.to_onnx_text(m))

create_model()