onnxruntime_extensions 0.14.0__cp313-cp313-macosx_11_0_arm64.whl

onnxruntime_extensions/pnp/_torchext.py

@@ -0,0 +1,310 @@
+import onnx
+import torch
+import numpy as np
+from typing import Any
+from onnx import helper
+from onnx import onnx_pb as onnx_proto
+from distutils.version import LooseVersion
+from torch.onnx import register_custom_op_symbolic
+
+from ._utils import ONNXModelUtils
+from ._base import CustomFunction, ProcessingTracedModule, is_processing_module
+from ._onnx_ops import ox as _ox, schema as _schema
+from ._onnx_ops import ONNXElementContainer, make_model_ex
+from .._ortapi2 import OrtPyFunction, get_opset_version_from_ort
+
+
+def _is_numpy_object(x):
+    return isinstance(x, (np.ndarray, np.generic))
+
+
+def _is_numpy_string_type(arr):
+    return arr.dtype.kind in {'U', 'S'}
+
+
+def _is_string_type(x):
+    if isinstance(x, list):
+        return any(_is_string_type(e) for e in x)
+    elif isinstance(x, torch.Tensor):
+        return False
+    elif not _is_numpy_object(x):
+        x = np.array(x)
+    return _is_numpy_string_type(x)
+
+
+def _to_onnx_type(dtype):
+    ty_dict = {torch.bool: onnx_proto.TensorProto.BOOL,
+               torch.float32: onnx_proto.TensorProto.FLOAT,
+               torch.float64: onnx_proto.TensorProto.DOUBLE,
+               torch.long: onnx_proto.TensorProto.INT64,
+               torch.int32: onnx_proto.TensorProto.INT32}
+    # ...
+    return ty_dict.get(dtype, onnx_proto.TensorProto.STRING)
+
+
+class OnnxOpFunction(CustomFunction):
+    @classmethod
+    def get_next_id_name(cls, name_base):
+        name = 'cls' if name_base is None else name_base
+        _cid = getattr(cls, '_cid', 1)
+        cls._cid = _cid + 1
+        return "{}_{}".format(name, _cid)
+
+    @staticmethod
+    def jvp(ctx: Any, *grad_inputs: Any) -> Any:
+        pass
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        return grad_outputs
+
+    @classmethod
+    def build_model(cls, opset_version, *args):
+        # build the one node graph
+        if isinstance(args[0], list):
+            args = [np.asarray(_i) for _i in args]
+        ec = ONNXElementContainer(get_opset_version_from_ort() if opset_version is None else opset_version)
+        attrs = cls.attrs
+        vi_inputs = [helper.make_tensor_value_info(
+            'it_' + str(id(_arg)), _to_onnx_type(_arg.dtype), list(_arg.shape))
+            for _arg in args]
+        inputs = [_vi.name for _vi in vi_inputs]
+        if hasattr(cls.opb_func, 'outputs') and len(cls.opb_func.outputs) > 0:
+            vi_outputs = [helper.make_tensor_value_info(
+                cls.get_next_id_name('ot'), *_schm) for _schm in cls.opb_func.outputs]
+        else:
+            vi_outputs = [helper.make_tensor_value_info(
+                cls.get_next_id_name('ot'), onnx_proto.TensorProto.FLOAT, []
+            )]
+        outputs = [_vi.name for _vi in vi_outputs]
+        # build the node
+        opfunc = cls.opb_func
+        opfunc(inputs, outputs, ec, None, **attrs)
+        g = helper.make_graph(ec.nodes, cls.get_next_id_name('g'), vi_inputs, vi_outputs)
+        m = make_model_ex(g, ec.node_domain_version_pair_sets, ec.target_opset)
+        return m
+
+    @classmethod
+    @torch.jit.unused
+    def _onnx_call(cls, ctx, *args) -> Any:
+        m = cls.build_model(None, *args)
+        try:
+            f = OrtPyFunction.from_model(m)
+            result = f(*list(_i.numpy() if isinstance(_i, torch.Tensor) else _i for _i in args))
+        except Exception as e:
+            onnx.save_model(m, '_temp_debugging.onnx')
+            raise e
+
+        results = result if isinstance(result, tuple) else [result]
+        return tuple([torch.from_numpy(_o) for _o in results]) if len(results) > 1 else torch.from_numpy(results[0])
+
+    @classmethod
+    def forward(cls, ctx: Any, *args: Any, **kwargs: Any) -> Any:
+        return cls._onnx_call(ctx, *args, **kwargs)
+
+    @classmethod
+    def symbolic(cls, g, *args):
+        return g.op(cls.op_type, *args)
+
+
+def create_op_function(op_type: str, func, **attrs):
+    if _ox.is_raw(func):
+        func = _schema(func.__func__)
+    cls = type(_ox.get_unique_operator_type_name(op_type), (OnnxOpFunction,),
+               dict(
+                   op_type=op_type,
+                   opb_func=func,
+                   attrs=attrs
+               ))
+    return cls.apply  # noqa
+
+
+onnx_pad = create_op_function('Pad', _ox.pad)
+onnx_where = create_op_function('Where', _ox.where)
+onnx_greater = create_op_function('Greater', _ox.greater)
+
+
+class _OnnxModelFunction:
+    id_object_map = {}  # cannot use the string directly since jit.script doesn't support the data type
+    id_function_map = {}
+    str_model_function_id = '_model_function_id'
+    str_model_id = '_model_id'
+    str_model_attached = '_model_attached'
+
+
+@torch.jit.ignore
+def _invoke_onnx_model(model_id: int, *args, **kwargs):
+    func = _OnnxModelFunction.id_function_map.get(model_id, None)
+    if not func:
+        model_or_path = _OnnxModelFunction.id_object_map.get(model_id)
+        if model_or_path is None:
+            raise ValueError("cannot find id={} registered!".format(model_id))
+        func = OrtPyFunction.from_model(model_or_path)
+        _OnnxModelFunction.id_function_map[model_id] = func
+    results = func(*list(_i.numpy() if isinstance(_i, torch.Tensor) else _i for _i in args), **kwargs)
+    return tuple(
+        [torch.from_numpy(_o) for _o in results]) if isinstance(results, tuple) else torch.from_numpy(results)
+
+
+@torch.jit.ignore
+def invoke_onnx_model1(model_id: int, arg0):
+    return _invoke_onnx_model(model_id, arg0)
+
+
+@torch.jit.ignore
+def invoke_onnx_model2(model_id: int, arg0, arg1):
+    return _invoke_onnx_model(model_id, arg0, arg1)
+
+
+@torch.jit.ignore
+def invoke_onnx_model3(model_id: int, arg0, arg1, arg2):
+    return _invoke_onnx_model(model_id, arg0, arg1, arg2)
+
+
+class _OnnxTracedFunction(CustomFunction):
+    @classmethod
+    def forward(cls, ctx: Any, *args: Any, **kwargs: Any) -> Any:
+        return _invoke_onnx_model(args[0].item(), *args[1:], **kwargs)
+
+    @classmethod
+    def symbolic(cls, g, *args):
+        ret = g.op('ai.onnx.contrib::_ModelFunctionCall', *args)
+        model_id = torch.onnx.symbolic_helper._maybe_get_scalar(args[0])  # noqa
+        if not model_id:
+            return ret
+
+        func = _OnnxModelFunction.id_function_map.get(model_id.item(), None)
+        if not func or len(func.outputs) <= 1:
+            return ret
+
+        outputs = [ret]
+        for _ in range(len(func.outputs) - 1):
+            outputs.append(ret.node().addOutput())
+
+        return tuple(outputs)
+
+
+def create_model_function(model_or_path):
+    _id = id(model_or_path)
+    assert _id != 0, "internal error: the id of a Python object is 0."
+    _OnnxModelFunction.id_object_map[_id] = model_or_path
+    return _id
+
+
+def get_id_models():
+    return _OnnxModelFunction.id_object_map
+
+
+class OnnxTracedModelFunction:
+    def __init__(self, onnx_model):
+        self.func_id = create_model_function(onnx_model)
+
+    def __call__(self, *args, **kwargs):
+        return _OnnxTracedFunction.apply(torch.tensor(self.func_id), *args, **kwargs)
+
+
+class _OnnxModelModule(torch.nn.Module):
+    def __init__(self, mdl):
+        super(_OnnxModelModule, self).__init__()
+        self.function = OnnxTracedModelFunction(mdl)
+
+    def forward(self, *args):
+        return self.function(*args)
+
+
+def _symbolic_pythonop(g: torch._C.Graph, n: torch._C.Node, *args, **kwargs):
+    name = kwargs["name"]
+    if name.startswith(invoke_onnx_model1.__name__[:-1]):
+        # NB: if you want to get the value of the first argument, i.e. the model id,
+        # you can get it by torch.onnx.symbolic_helper._maybe_get_scalar(args[0]).item()
+        ret = g.op("ai.onnx.contrib::_ModelFunctionCall", *args)
+    else:
+        # Logs a warning and returns None
+        import warnings
+        return warnings.warn("prim::PythonOp", "unknown node kind: " + name)
+    # Copy type and shape from original node.
+    ret.setType(n.output().type())
+    return ret
+
+
+if LooseVersion(torch.__version__) >= LooseVersion("1.11"):
+    register_custom_op_symbolic("prim::PythonOp", _symbolic_pythonop, 1)
+
+
+class SequentialProcessingModule(ProcessingTracedModule):
+    def __init__(self, *models):
+        super(SequentialProcessingModule, self).__init__()
+        self.model_list = torch.nn.ModuleList()
+        for mdl_ in models:
+            if isinstance(mdl_, onnx.ModelProto):
+                self.model_list.append(_OnnxModelModule(mdl_))
+            elif is_processing_module(mdl_):
+                self.model_list.append(mdl_)
+            else:
+                assert callable(mdl_), "the model type is not recognizable."
+                self.model_list.append(ProcessingTracedModule(mdl_))
+
+    def forward(self, *args):
+        outputs = args
+        with torch.no_grad():
+            for idx_, mdl_ in enumerate(self.model_list):
+                if not isinstance(outputs, tuple):
+                    outputs = (outputs,)
+                outputs = mdl_(*outputs)
+
+        return outputs
+
+    def export(self, *args, **kwargs):
+        prefix_m = None
+        core_m = self
+        raw_input_flag = any(_is_string_type(x_) for x_ in args)
+        if raw_input_flag:
+            # NB: torch.onnx.export doesn't support exporting a module accepting string type input,
+            # So, in this case, the module will be separated into two parts to use the customized export.
+            m0 = self.model_list[0]
+            new_args = m0(*args)
+            if not isinstance(new_args, tuple):
+                new_args = (new_args, )
+            prefix_m = m0.export(*args, **kwargs)
+            args = new_args
+            core_m = SequentialProcessingModule(*self.model_list[1:])
+        if prefix_m is None:
+            return super().export(*args, **kwargs)
+        else:
+            oxml = core_m.export(*args, **kwargs)
+            model = ONNXModelUtils.join_models(prefix_m, oxml)
+
+            # Rename the input/output node names if the user has provided any substitutions!
+            # Ref: https://github.com/onnx/onnx/issues/2052
+            # Known issue: This logic doesn't deal with subgraphs.
+            if (('input_names' in kwargs) or ('output_names' in kwargs)) and \
+              (kwargs['input_names'] or kwargs['output_names']):
+                swaps = {}
+                if 'input_names' in kwargs and kwargs['input_names']:
+                    assert len(model.graph.input) == len(kwargs['input_names']), \
+                      "Expecting {} input names but got {}".format(
+                        len(model.graph.input), len(kwargs['input_names']))
+                    for n, new_name in zip(model.graph.input, kwargs['input_names']):
+                        swaps[n.name] = new_name
+                        n.name = new_name
+
+                if 'output_names' in kwargs and kwargs['output_names']:
+                    assert len(model.graph.output) == len(kwargs['output_names']), \
+                      "Expecting {} output names but got {}".format(
+                        len(model.graph.output), len(kwargs['output_names']))
+                    for n, new_name in zip(model.graph.output, kwargs['output_names']):
+                        swaps[n.name] = new_name
+                        n.name = new_name
+
+                if swaps:
+                    for n in model.graph.node:
+                        for j in range(len(n.input)):
+                            n.input[j] = swaps.get(n.input[j], n.input[j])
+
+                        for j in range(len(n.output)):
+                            n.output[j] = swaps.get(n.output[j], n.output[j])
+
+                    for n in model.graph.initializer:
+                        n.name = swaps.get(n.name, n.name)
+
+            return model

onnxruntime_extensions/pnp/_unifier.py

@@ -0,0 +1,45 @@
+import onnx
+
+from .._ortapi2 import get_opset_version_from_ort
+from ._utils import ONNXModelUtils
+from ._base import is_processing_module
+from ._torchext import get_id_models, SequentialProcessingModule
+
+
+def export(m, *args,
+           opset_version=0,
+           output_path=None,
+           export_params=True,
+           verbose=False,
+           input_names=None,
+           output_names=None,
+           operator_export_type=None,
+           do_constant_folding=True,
+           dynamic_axes=None,
+           keep_initializers_as_inputs=None,
+           custom_opsets=None,
+           io_mapping=None):
+    """
+    export all models and modules into a merged ONNX model.
+    """
+    if opset_version == 0:
+        opset_version = get_opset_version_from_ort()
+
+    if not is_processing_module(m):
+        m = SequentialProcessingModule(m)
+
+    model = m.export(*args, opset_version=opset_version,
+                     output_path=output_path,
+                     export_params=export_params,
+                     verbose=verbose,
+                     input_names=input_names,
+                     output_names=output_names,
+                     operator_export_type=operator_export_type,
+                     do_constant_folding=do_constant_folding,
+                     dynamic_axes=dynamic_axes,
+                     keep_initializers_as_inputs=keep_initializers_as_inputs,
+                     custom_opsets=custom_opsets)
+    full_m = ONNXModelUtils.unfold_model(model, get_id_models(), io_mapping)
+    if output_path is not None:
+        onnx.save_model(full_m, output_path)
+    return full_m

onnxruntime_extensions/pnp/_utils.py

@@ -0,0 +1,302 @@
+import copy
+import onnx
+from onnx import helper, numpy_helper
+from collections import namedtuple
+
+
+class _Container:
+    def __init__(self):
+        self.parent = None
+        self.initializer=[]
+        self.value_info=[]
+        self.nodes = []
+        self.node_domain_version_pair_sets = {}
+
+    def add_model(self, oxml):
+        self.initializer.extend(oxml.graph.initializer)
+        self.value_info.extend(oxml.graph.value_info)
+        self.nodes.extend(oxml.graph.node)
+        self.node_domain_version_pair_sets.update(
+            [(opset_.domain, opset_.version) for opset_ in oxml.opset_import])
+        return self
+
+
+class ONNXModelUtils:
+    @staticmethod
+    def merge_name(prefix, name):
+        return "{}_{}".format(prefix, name)
+
+    @staticmethod
+    def _rename_iter(iterables, prefix_name, inplace=False):
+        new_iz = iterables if inplace else [copy.deepcopy(iz_) for iz_ in iterables]
+        for iz_ in new_iz:
+            iz_.name = ONNXModelUtils.merge_name(prefix_name, iz_.name)
+        return new_iz
+
+    @classmethod
+    def _rename_graph(cls, graph, prefix, graph_or_container):
+        def io_rename(node, prefix_name, idx):
+            new_node = copy.deepcopy(node)
+            if not node.name:
+                new_node.name = cls.merge_name(prefix_name, "op{}".format(idx))
+            else:
+                new_node.name = cls.merge_name(prefix_name, node.name)
+
+            del new_node.input[:]
+            new_node.input.extend(ONNXModelUtils.merge_name(prefix_name, nm_) if nm_ else '' for nm_ in node.input)
+            del new_node.output[:]
+            new_node.output.extend(ONNXModelUtils.merge_name(prefix_name, nm_) if nm_ else '' for nm_ in node.output)
+            return new_node
+
+        assert prefix is not None, 'The graph prefix could not be None'
+        graph_or_container.initializer.extend(cls._rename_iter(graph.initializer, prefix))
+        graph_or_container.value_info.extend(cls._rename_iter(graph.value_info, prefix))
+        return list(io_rename(nd_, prefix, idx_) for idx_, nd_ in enumerate(graph.node))
+
+    @classmethod
+    def _process_node_body(cls, node, prefix):
+        if all(attr.name != 'body' for attr in node.attribute):
+            return node
+
+        def _process_attr(attr, prefix_name):
+            if attr.name == 'body':
+                new_attr = copy.deepcopy(attr)
+                del new_attr.g.value_info[:]
+                del new_attr.g.node[:]
+                new_attr.g.node.extend(cls._rename_graph(attr.g, prefix_name, new_attr.g))
+                cls._rename_iter(new_attr.g.input, prefix_name, inplace=True)
+                cls._rename_iter(new_attr.g.output, prefix_name, inplace=True)
+                return new_attr
+            else:
+                return attr
+
+        attr_list = list(_process_attr(attr_, prefix) for attr_ in node.attribute)
+        del node.attribute[:]
+        node.attribute.extend(attr_list)
+        return node
+
+    @staticmethod
+    def get_model_name_abbr(node):
+        no = node.name.split('_')[-1]
+        return 'm_' + no
+
+    @staticmethod
+    def get_model_id_from_arg0(nodes, node):
+        arg0_name = node.input[0]
+        c_node = [n_ for n_ in nodes if
+                  n_.op_type == 'Constant' and n_.output[0] == arg0_name]
+        assert len(c_node) == 1, 'internal error, multiple nodes with the same output.'
+        c_node = c_node[0]
+        tensor_value = onnx.helper.get_attribute_value(c_node.attribute[0])
+        _id = numpy_helper.to_array(tensor_value).item()
+        return _id
+
+    @classmethod
+    def _unfold_model_node(cls, container, name, model, io_mapping=None):
+        top_container = container
+        while top_container.parent is not None:  # only one opset_import in the model.
+            top_container = top_container.parent
+
+        renamed_nodes = cls._rename_graph(model.graph, name, container)
+        onnx_nodes = [cls._process_node_body(nd_, name) for nd_ in renamed_nodes]
+
+        top_container.node_domain_version_pair_sets.update(
+            [(opset_.domain, opset_.version) for opset_ in model.opset_import])
+        return onnx_nodes
+
+    @classmethod
+    def unfold_model(cls, oxml, id_to_model, io_mapping=None):
+        container = _Container().add_model(oxml)
+        nodes = []
+        for _nid, _node in enumerate(oxml.graph.node):
+            if _node.op_type != '_ModelFunctionCall':
+                nodes.append(_node)
+            else:
+                model_id = cls.get_model_id_from_arg0(list(oxml.graph.node), _node)
+                if model_id not in id_to_model:
+                    raise RuntimeError("Cannot find the model id({}) in the table".format(model_id))
+
+                prefix = cls.get_model_name_abbr(_node)
+                nest_model = id_to_model[model_id]
+
+                input_mapping = []
+                output_mapping = []
+                for idx_, in_ in enumerate(nest_model.graph.input):
+                    _renamed_in = "{}_{}".format(prefix, in_.name)
+                    _nd = onnx.helper.make_node('Identity',
+                                                [_node.input[idx_ + 1]],  # the first arg is model id, skip it.
+                                                [_renamed_in],
+                                                name='i_' + _renamed_in)
+                    input_mapping.append(_nd)
+                nds = cls._unfold_model_node(container,
+                                             prefix,
+                                             nest_model,
+                                             io_mapping)
+                for idx_, out_ in enumerate(nest_model.graph.output):
+                    if idx_ >= len(_node.output):
+                        continue
+                    _renamed_out = "{}_{}".format(prefix, out_.name)
+                    _nd = onnx.helper.make_node('Identity',
+                                                [_renamed_out],
+                                                [_node.output[idx_]],
+                                                name='o_' + _renamed_out)
+                    output_mapping.append(_nd)
+                if io_mapping is not None:
+                    assert callable(io_mapping), "io_mapping is a custom function to build the linkage of the models"
+                    input_mapping, output_mapping = io_mapping(input_mapping, output_mapping)
+                # attention: the order of the list operations is important, which avoids the topological sort.
+                nodes.extend(input_mapping)
+                nodes.extend(nds)
+                nodes.extend(output_mapping)
+
+        intlzs = cls._remove_unused_initializers(nodes, container.initializer)
+        oxml = copy.deepcopy(oxml)
+        del oxml.graph.node[:]
+        oxml.graph.node.extend(nodes)
+        del oxml.graph.initializer[:]
+        oxml.graph.initializer.extend(intlzs)
+        return oxml
+
+    @classmethod
+    def topological_sort(cls, container, nodes, inputs, outputs):
+        op_output_map = {}
+        DynNode = namedtuple('DynNode', ['name', 'output'])
+        input_nodes = [DynNode(name='placeholder',
+                               output=[nm_.name for nm_ in inputs] +
+                                      [it_.name for it_ in container.initializers])] + \
+                      [nd_ for nd_ in nodes if nd_.op_type == 'Constant']
+
+        for nd_ in nodes + input_nodes:
+            for ky_ in nd_.output:
+                op_output_map[ky_] = nd_
+
+        edges = {}
+        for op in nodes:
+            for x in op.input:
+                if x == '':
+                    continue
+                try:
+                    predecessor = op_output_map[x]
+                except KeyError:
+                    raise RuntimeError(
+                        "{}: cannot find an operator to produce the tensor: {}".format(op.name, x)) from None
+
+                val = edges.get(predecessor.name, [])
+                val.append(op)
+                edges[predecessor.name] = val
+
+        for y_ in outputs:
+            op = op_output_map[y_.name].name
+            if op not in edges:
+                edges[op] = []
+
+        visited = set()
+        sorted_nodes = []
+        unfinished_nodes = set()
+
+        def recursive_helper(node):
+            if node.name in visited:
+                return
+
+            if node.name in unfinished_nodes:
+                raise RuntimeError("ONNX Graph is not a DAG, the cycle is found at {}".format(node.name))
+
+            unfinished_nodes.add(node.name)
+            if node.name in edges:  # if the node's output is not in the Graph output.
+                assert node.name != '', 'this topological-sort depends on the unique node name.'
+                for successor in edges[node.name]:
+                    recursive_helper(successor)
+
+            unfinished_nodes.remove(node.name)
+            visited.add(node.name)
+            if node is not input_nodes[0]:
+                sorted_nodes.insert(0, node)
+
+        for nd_ in input_nodes:
+            recursive_helper(nd_)
+
+        return sorted_nodes
+
+    @staticmethod
+    def _remove_unused_initializers(nodes, initializers, reserved_names=None):
+        if reserved_names is None:
+            reserved_names = set()
+        nodes_input_set = set()
+        for nd_ in nodes:
+            nodes_input_set.update(n_ for n_ in nd_.input)
+
+        return [intlz_ for intlz_ in initializers if intlz_.name in nodes_input_set or intlz_.name in reserved_names]
+
+    @classmethod
+    def join_models(cls, *models, io_mapping=None):
+        # generate the prefix id for the embedding graph to avoid the name conflict
+        mdl_prefix = []
+        for _i in range(len(models)):
+            mdl_prefix.append("g{}".format(_i + 1))
+
+        inputs = cls._rename_iter(models[0].graph.input, mdl_prefix[0])
+        outputs = cls._rename_iter(models[-1].graph.output, mdl_prefix[-1])
+
+        port_mapping = {}
+        if io_mapping is not None:
+            assert callable(io_mapping), "io_mapping is a custom function to build the linkage of the models"
+            ModelPort = namedtuple('ModelPort', "input output")
+            ports = []
+            for _idx in range(len(models)):
+                mio = ModelPort([cls.merge_name(mdl_prefix[_idx], _x.name) for _x in models[_idx].graph.input],
+                                [cls.merge_name(mdl_prefix[_idx], _y.name) for _y in models[_idx].graph.output])
+                ports.append(mio)
+            port_mapping = io_mapping(ports)
+        for _idx in range(len(models) - 1):
+            for _i, _x in enumerate(models[_idx + 1].graph.input):
+                iname = cls.merge_name(mdl_prefix[_idx + 1], _x.name)
+                if iname not in port_mapping:
+                    oname = cls.merge_name(mdl_prefix[_idx], models[_idx].graph.output[_i].name)
+                    port_mapping[iname] = oname
+
+        nodes = []
+        container = _Container()
+        for _idx, _m in enumerate(models):
+            container.add_model(_m)
+            nodes += cls._rename_graph(_m.graph, mdl_prefix[_idx], container)
+
+        for _n in nodes:
+            replaceable = False
+            for _i in _n.input:
+                if _i in port_mapping:
+                    replaceable = True
+                    break
+            if replaceable:
+                new_input = copy.deepcopy(_n.input)
+                del _n.input[:]
+                _n.input.extend([port_mapping[_i] if _i in port_mapping else _i for _i in new_input])
+
+        name = "_".join([_mdl.graph.name for _mdl in models])
+        domains = set()
+        _opset = []
+        for _mdl in models:
+            for _ops in _mdl.opset_import:
+                domain = _ops.domain if _ops.domain else "ai.onnx"
+                if domain in domains:
+                    if domain == "ai.onnx":
+                      assert _ops.version == _opset[0].version, \
+                        f"ai.onnx domain version doesn't match {_ops.version} != {_opset[0].version}"
+                else:
+                    domains.add(domain)
+                    if domain == "ai.onnx":
+                        _opset.insert(0, _ops)
+                    else:
+                        _opset.append(_ops)
+
+        inits = cls._remove_unused_initializers(nodes, container.initializer)
+        g = helper.make_graph(nodes, name, inputs, outputs,
+                              initializer=inits,
+                              value_info=container.value_info)
+
+        if hasattr(helper, 'make_model_gen_version'):
+            # make_model_gen_version doesn't accept the custom domain.
+            m = helper.make_model_gen_version(g, opset_imports=_opset[:1])
+            m.opset_import.extend(_opset[1:])
+        else:
+            m = helper.make_model(g, opset_imports=_opset)
+        return m