code: def _add_onnxscript_fn(
      model_bytes: bytes,
      custom_opsets: Mapping[str, int],
  ) -> bytes:
      """Insert model-included custom onnx-script function into ModelProto"""
      # TODO(titaiwang): remove this when onnx becomes dependency
      try:
          import onnx
      except ImportError as e:
          raise errors.OnnxExporterError(
              "Module onnx is not installed!"
          ) from e
  
      # For > 2GB model, onnx.load_fromstring would fail. However, because
      # in _export_onnx, the tensors should be saved separately if the proto
      # size > 2GB, and if it for some reason did not, the model would fail on
      # serialization anyway in terms of the protobuf limitation. So we don't
      # need to worry about > 2GB model getting here.
      model_proto = onnx.load_model_from_string(model_bytes)  # type: ignore[attr-defined]
  
      # Iterate graph nodes to insert only the included custom
      # function_proto into model_proto
      # TODO(titaiwang): Currently, onnxscript doesn't support ONNXFunction
      # calling other ONNXFunction scenario, neither does it here
      onnx_function_list = list()  # type: ignore[var-annotated]
      included_node_func = set()  # type: Set[str]
      # onnx_function_list and included_node_func are expanded in-place
      _find_onnxscript_op(
          model_proto.graph,
          included_node_func,
          custom_opsets,
          onnx_function_list,
      )
  
      if onnx_function_list:
          model_proto.functions.extend(onnx_function_list)
          model_bytes = model_proto.SerializeToString()
      return model_bytes
  
file_path: pytorch/torch/onnx/_internal/onnx_proto_utils.py
imports: from __future__ import annotations
  import glob
  import io
  import os
  import shutil
  import zipfile
  from typing import Any, List, Mapping, Set, Tuple, Union
  import torch
  import torch.jit._trace
  import torch.serialization
  from torch.onnx import _constants, _exporter_states, errors
  from torch.onnx._internal import _beartype, jit_utils, registration
label_desc: Use image node_id 5 for calling a global function with example usage: _add_onnxscript_fn(model_bytes, custom_opsets) and returns: model_bytes