onnx 1.18.0__cp310-cp310-win32.whl → 1.19.1rc1__cp310-cp310-win32.whl

onnx/reference/reference_evaluator.py

@@ -8,10 +8,8 @@ from typing import Any
 
 import numpy as np
 
-from onnx import load
-from onnx.defs import onnx_opset_version
-from onnx.external_data_helper import ExternalDataInfo, uses_external_data
-from onnx.model_container import ModelContainer
+import onnx
+import onnx.model_container
 from onnx.onnx_pb import (
     FunctionProto,
     GraphProto,
@@ -20,13 +18,7 @@ from onnx.onnx_pb import (
     TensorProto,
     TypeProto,
 )
-from onnx.reference.op_run import (
-    OpFunctionContextDependant,
-    OpRun,
-    OpRunExpand,
-    RuntimeContextError,
-    to_array_extended,
-)
+from onnx.reference import op_run
 from onnx.reference.ops_optimized import optimized_operators
 
 
@@ -114,7 +106,7 @@ class ReferenceEvaluator:
 
             op_domain = "custom"
 
-            def _run(self, x, alpha=None):  # type: ignore
+            def _run(self, x, alpha=None):
                 # None must be the default value, it is automatically
                 # replaced by class OpRun with either the default value
                 # specified in the NodeProto or an attribute value defined
@@ -204,9 +196,9 @@ class ReferenceEvaluator:
         self,
         proto: Any,
         opsets: dict[str, int] | None = None,
-        functions: list[ReferenceEvaluator | FunctionProto] | None = None,  # type: ignore
+        functions: list[ReferenceEvaluator | FunctionProto] | None = None,
         verbose: int = 0,
-        new_ops: list[type[OpRun]] | None = None,
+        new_ops: list[type[op_run.OpRun]] | None = None,
         optimized: bool = True,
     ) -> None:
         if optimized:
@@ -220,17 +212,17 @@ class ReferenceEvaluator:
         self.output_types_ = None
         self.input_types_ = None
 
-        if isinstance(proto, ModelContainer):
-            self.container_: ModelContainer | None = proto
+        if isinstance(proto, onnx.model_container.ModelContainer):
+            self.container_: onnx.model_container.ModelContainer | None = proto
             proto = self.container_.model_proto
         else:
             self.container_ = None
 
         if isinstance(proto, str):
             with open(proto, "rb") as f:
-                proto = load(f)
+                proto = onnx.load(f)
         elif isinstance(proto, bytes):
-            proto = load(BytesIO(proto))
+            proto = onnx.load(BytesIO(proto))
         self.proto_ = proto
         self.functions_: dict[tuple[str, str], ReferenceEvaluator] = {}
         self.attributes_: list[str] = []
@@ -248,15 +240,17 @@ class ReferenceEvaluator:
                 raise TypeError("opsets must be a dictionary if proto is GraphProto.")
             self.opsets_ = opsets
         elif isinstance(proto, FunctionProto):
-            self.onnx_graph_ = None  # type: ignore
+            self.onnx_graph_ = None
             self.opsets_ = {d.domain: d.version for d in proto.opset_import}
             if opsets is not None:
                 raise ValueError("opsets must be None if proto is FunctionProto.")
             self.attributes_ = list(proto.attribute)
         elif isinstance(proto, NodeProto):
-            self.onnx_graph_ = None  # type: ignore
+            self.onnx_graph_ = None
             self.opsets_ = {
-                proto.domain: 1 if proto.domain != "" else onnx_opset_version()
+                proto.domain: 1
+                if proto.domain != ""
+                else onnx.defs.onnx_opset_version()
             }
         else:
             raise TypeError(f"Unexpected type {type(proto)} for proto.")
@@ -266,7 +260,7 @@ class ReferenceEvaluator:
             self.output_names_ = [o.name for o in self.onnx_graph_.output]
             self.output_types_ = [i.type for i in self.onnx_graph_.output]
             self.inits_ = list(self.onnx_graph_.initializer) + list(
-                self.onnx_graph_.sparse_initializer  # type: ignore
+                self.onnx_graph_.sparse_initializer
             )
             self.nodes_ = self.onnx_graph_.node
             all_types = {i.name: i.type for i in self.onnx_graph_.input}
@@ -283,27 +277,27 @@ class ReferenceEvaluator:
             else:
                 self.nodes_ = proto.node
         if functions is not None:
-            for f in functions:  # type: ignore
+            for f in functions:
                 if isinstance(f, FunctionProto):
                     self.functions_[f.domain, f.name] = self.__class__(
                         f, verbose=verbose, functions=list(self.functions_.values())
                     )
                 elif isinstance(f, ReferenceEvaluator):
-                    onx = f.proto_  # type: ignore
+                    onx = f.proto_
                     self.functions_[onx.domain, onx.name] = f
                 else:
                     raise TypeError(f"Unexpected type {type(f)!r} for a function.")
         self.verbose = verbose
-        self.new_ops_: dict[tuple[str, str], type[OpRun]] = {}
+        self.new_ops_: dict[tuple[str, str], type[op_run.OpRun]] = {}
         if new_ops is not None:
             for cl in new_ops:
                 if not hasattr(cl, "op_domain"):
                     raise AttributeError(
                         f"Class {cl} must define attribute 'op_domain'."
                     )
-                if not issubclass(cl, OpRun):  # type: ignore
+                if not issubclass(cl, op_run.OpRun):
                     raise TypeError(f"Class {cl} must inherit from OpRun (in new_ops).")
-                key = cl.op_domain, cl.__name__  # type: ignore
+                key = cl.op_domain, cl.__name__
                 if key in self.new_ops_:
                     # Already an implementation, the first one is used.
                     continue
@@ -312,7 +306,7 @@ class ReferenceEvaluator:
 
     def retrieve_external_data(self, initializer: TensorProto) -> np.ndarray:
         """Returns a tensor saved as external."""
-        info = ExternalDataInfo(initializer)
+        info = onnx.external_data_helper.ExternalDataInfo(initializer)
         location = info.location
         if self.container_ and self.container_.is_in_memory_external_initializer(
             location
@@ -349,27 +343,27 @@ class ReferenceEvaluator:
             print(pattern % tuple(new_args))
 
     @property
-    def input_names(self):  # type: ignore
+    def input_names(self):
         """Returns the input names."""
         return self.input_names_
 
     @property
-    def input_types(self):  # type: ignore
+    def input_types(self):
         """Returns the input types if any specified."""
         return self.input_types_
 
     @property
-    def output_names(self):  # type: ignore
+    def output_names(self):
         """Returns the output names."""
         return self.output_names_
 
     @property
-    def output_types(self):  # type: ignore
+    def output_types(self):
         """Returns the output types."""
         return self.output_types_
 
     @property
-    def opsets(self):  # type: ignore
+    def opsets(self):
         """Returns the opsets."""
         return self.opsets_
 
@@ -403,8 +397,8 @@ class ReferenceEvaluator:
         for init in self.inits_:
             self.rt_inits_[init.name] = (
                 self.retrieve_external_data(init)
-                if uses_external_data(init)
-                else to_array_extended(init)
+                if onnx.external_data_helper.uses_external_data(init)
+                else onnx.numpy_helper.to_array(init)
             )
         run_params = {
             "log": lambda pattern, *args: self._log(10, pattern, *args),
@@ -421,25 +415,26 @@ class ReferenceEvaluator:
                     all_types[shape_type.name] = shape_type.type
             self.all_types_ = all_types
         else:
-            self.all_types_ = None  # type: ignore
+            self.all_types_ = None
 
         for node in self.nodes_:
             try:
                 cl = self._load_impl(node)
-            except RuntimeContextError as e:
+            except op_run.RuntimeContextError as e:
                 # A node has a context dependent implementation.
                 # Shape inference must be run to get the input types.
                 if self.all_types_:
                     it = [self.get_result_types(i, exc=False) for i in node.input]
                     if None in it:
                         # One input does not exist. It must be done while executing the graph.
-                        cl = lambda *args, parent=self: OpFunctionContextDependant(  # noqa: E731
-                            *args, parent=parent
-                        )
+                        def cl(*args, parent=self):
+                            return op_run.OpFunctionContextDependant(
+                                *args, parent=parent
+                            )
                     else:
-                        cl = self._load_impl(node, it)  # type: ignore
+                        cl = self._load_impl(node, it)
                 else:
-                    raise RuntimeContextError(
+                    raise op_run.RuntimeContextError(
                         f"No implementation was found for node type {node.op_type!r} from domain {node.domain!r}. "
                         f"If this node has a context dependent implementation, you should run function infer_shapes "
                         f"before calling ReferenceEvaluator."
@@ -470,13 +465,13 @@ class ReferenceEvaluator:
             # This mechanism can be used to implement a custom onnx node
             # or to overwrite an existing one.
             cl = self.new_ops_[key]
-            if not issubclass(cl, OpRunExpand):
+            if not issubclass(cl, op_run.OpRunExpand):
                 return cl
             # It must be replaced by its implementation defined in its schema.
             expand = True
 
         if node.domain == "":
-            from onnx.reference.ops import load_op
+            from onnx.reference.ops import load_op  # noqa: PLC0415
 
             try:
                 return load_op(
@@ -486,7 +481,7 @@ class ReferenceEvaluator:
                     expand=expand,
                     evaluator_cls=self.__class__,
                 )
-            except RuntimeContextError:
+            except op_run.RuntimeContextError:
                 if input_types is None:
                     raise
                 return load_op(
@@ -506,23 +501,29 @@ class ReferenceEvaluator:
                 f"{node.domain},{node.op_type} from the list of inlined operator."
             )
         if node.domain == "ai.onnx.preview.training":
-            from onnx.reference.ops.aionnx_preview_training import load_op as load_op_pt
+            from onnx.reference.ops.aionnx_preview_training import (  # noqa: PLC0415
+                load_op as load_op_pt,
+            )
 
             return load_op_pt(node.domain, node.op_type, version)
 
         if node.domain == "experimental":
-            from onnx.reference.ops.experimental import load_op as load_op_exp
+            from onnx.reference.ops.experimental import (  # noqa: PLC0415
+                load_op as load_op_exp,
+            )
 
             return load_op_exp(node.domain, node.op_type, version)
 
         if node.domain == "ai.onnx.ml":
-            from onnx.reference.ops.aionnxml import load_op as load_op_ml
+            from onnx.reference.ops.aionnxml import (  # noqa: PLC0415
+                load_op as load_op_ml,
+            )
 
             return load_op_ml(node.domain, node.op_type, version)
 
         # It has to be a function.
         if key in self.functions_:
-            from onnx.reference.ops import load_op
+            from onnx.reference.ops import load_op  # noqa: PLC0415
 
             impl = self.functions_[key]
             return load_op(
@@ -543,7 +544,7 @@ class ReferenceEvaluator:
         feed_inputs: dict[str, Any],
         attributes: dict[str, Any] | None = None,
         intermediate: bool = False,
-    ) -> dict[str, Any] | list[Any]:  # type: ignore
+    ) -> dict[str, Any] | list[Any]:
         """Executes the onnx model.
 
         Args:

onnx/serialization.py

@@ -57,7 +57,7 @@ class _Registry:
     def register(self, serializer: ProtoSerializer) -> None:
         self._serializers[serializer.supported_format] = serializer
         self._extension_to_format.update(
-            {ext: serializer.supported_format for ext in serializer.file_extensions}
+            dict.fromkeys(serializer.file_extensions, serializer.supported_format)
         )
 
     def get(self, fmt: str) -> ProtoSerializer:
@@ -118,7 +118,7 @@ class _ProtobufSerializer(ProtoSerializer):
             raise TypeError(
                 f"Parameter 'serialized' must be bytes, but got type: {type(serialized)}"
             )
-        decoded = typing.cast(Optional[int], proto.ParseFromString(serialized))
+        decoded = typing.cast("Optional[int]", proto.ParseFromString(serialized))
         if decoded is not None and decoded != len(serialized):
             raise google.protobuf.message.DecodeError(
                 f"Protobuf decoding consumed too few bytes: {decoded} out of {len(serialized)}"
@@ -130,7 +130,7 @@ class _TextProtoSerializer(ProtoSerializer):
     """Serialize and deserialize text proto."""
 
     supported_format = "textproto"
-    file_extensions = frozenset({".textproto", ".prototxt", ".pbtxt"})
+    file_extensions = frozenset({".txtpb", ".textproto", ".prototxt", ".pbtxt"})
 
     def serialize_proto(self, proto: _Proto) -> bytes:
         textproto = google.protobuf.text_format.MessageToString(proto)

onnx/shape_inference/implementation.cc

@@ -710,7 +710,7 @@ class ShapeInferenceImplBase {
   void FinalizeShapeInference() {
     auto& errors = getErrors();
     // Throw shape inference error if any. Error mode right now only supports 0 and 1.
-    // When set to 0, any node level shape inference errors are not thrown. This is to support backward compatiblity
+    // When set to 0, any node level shape inference errors are not thrown. This is to support backward compatibility
     // with 1.7 and earlier releases. When set to 1 it will throw all exceptions.
     // TODO: Add a more granular way for exception handling.
     if (!errors.empty() && options.error_mode > 0) {

onnx/shape_inference.py

@@ -14,11 +14,14 @@ from typing import TYPE_CHECKING
 
 import onnx
 import onnx.onnx_cpp2py_export.shape_inference as C  # noqa: N812
-from onnx import AttributeProto, FunctionProto, ModelProto, TypeProto
+from onnx.onnx_pb import AttributeProto, FunctionProto, ModelProto, TypeProto
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+GraphInferencer = C.GraphInferencer
+InferenceContext = C.InferenceContext
+
 
 def infer_shapes(
     model: ModelProto | bytes,

onnx/subbyte.py

@@ -22,13 +22,7 @@ UINT4_MAX = 15
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32_to_4bit_unpacked(*args, **kwargs):
-    return _float32_to_4bit_unpacked(*args, **kwargs)
-
-
-def _float32_to_4bit_unpacked(
-    x: np.ndarray | np.dtype | float, signed: bool
-) -> np.ndarray:
+def float32_to_4bit_unpacked(x: np.ndarray | float, signed: bool) -> np.ndarray:
     """Cast to 4bit via rounding and clipping (without packing).
 
     Args:
@@ -52,11 +46,7 @@ def _float32_to_4bit_unpacked(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32x2_to_4bitx2(*args, **kwargs):
-    return _float32x2_to_4bitx2(*args, **kwargs)
-
-
-def _float32x2_to_4bitx2(
+def float32x2_to_4bitx2(
     val_low: np.dtype, val_high: np.dtype, signed: bool
 ) -> np.ndarray:
     """Cast two elements to 4bit (via rounding and clipping) and pack
@@ -69,8 +59,8 @@ def _float32x2_to_4bitx2(
     Returns:
         An ndarray with a single int8/uint8 element, containing both int4 elements
     """
-    i8_high = _float32_to_4bit_unpacked(val_high, signed)
-    i8_low = _float32_to_4bit_unpacked(val_low, signed)
+    i8_high = float32_to_4bit_unpacked(val_high, signed)  # type: ignore[arg-type]
+    i8_low = float32_to_4bit_unpacked(val_low, signed)  # type: ignore[arg-type]
     return i8_high << 4 | i8_low & 0x0F
 
 
@@ -78,11 +68,7 @@ def _float32x2_to_4bitx2(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def unpack_4bitx2(*args, **kwargs):
-    return _unpack_4bitx2(*args, **kwargs)
-
-
-def _unpack_4bitx2(
+def unpack_4bitx2(
     x: npt.NDArray[np.uint8], dims: int | Sequence[int]
 ) -> npt.NDArray[np.uint8]:
     """Unpack an array of packed uint8 elements (4bitx2) into individual elements
@@ -113,11 +99,7 @@ def _unpack_4bitx2(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def unpack_single_4bitx2(*args, **kwargs):
-    return _unpack_single_4bitx2(*args, **kwargs)
-
-
-def _unpack_single_4bitx2(
+def unpack_single_4bitx2(
     x: np.ndarray | np.dtype | float, signed: bool
 ) -> tuple[np.ndarray, np.ndarray]:
     def unpack_signed(x):
@@ -144,11 +126,7 @@ def _unpack_single_4bitx2(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32_to_float4e2m1_unpacked(*args, **kwargs):
-    return _float32_to_float4e2m1_unpacked(*args, **kwargs)
-
-
-def _float32_to_float4e2m1_unpacked(values: np.ndarray) -> np.ndarray:
+def float32_to_float4e2m1_unpacked(values: np.ndarray) -> np.ndarray:
     """Cast float32 to float4e2m1 (without packing).
 
     Args:
@@ -176,11 +154,7 @@ def _float32_to_float4e2m1_unpacked(values: np.ndarray) -> np.ndarray:
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32x2_to_float4e2m1x2(*args, **kwargs):
-    return _float32x2_to_float4e2m1x2(*args, **kwargs)
-
-
-def _float32x2_to_float4e2m1x2(val_low: np.ndarray, val_high: np.ndarray) -> np.ndarray:
+def float32x2_to_float4e2m1x2(val_low: np.ndarray, val_high: np.ndarray) -> np.ndarray:
     """Cast two elements to float4e2m1 and pack to a single byte
     Args:
         val_low: element to be packed in the 4 LSB
@@ -189,6 +163,6 @@ def _float32x2_to_float4e2m1x2(val_low: np.ndarray, val_high: np.ndarray) -> np.
     Returns:
         An ndarray with uint8 elements, containing both float4e2m1 elements
     """
-    i8_high = _float32_to_float4e2m1_unpacked(val_high)
-    i8_low = _float32_to_float4e2m1_unpacked(val_low)
+    i8_high = float32_to_float4e2m1_unpacked(val_high)
+    i8_low = float32_to_float4e2m1_unpacked(val_low)
     return i8_high << 4 | i8_low & 0x0F

onnx/test/basic_test.py

@@ -227,6 +227,90 @@ class TestBasicFunctions(unittest.TestCase):
         # Check if the version is correct.
         self.assertEqual(model.ir_version, onnx.IR_VERSION)
 
+    def test_model_and_graph_repr(self) -> None:
+        # Check if the __repr__ methods work without error
+        model = _simple_model()
+        model_repr = repr(model)
+        self.assertEqual(
+            model_repr,
+            "ModelProto(ir_version=12, producer_name='onnx-test', graph=GraphProto('test'))",
+        )
+
+        text_model = """
+           <
+             ir_version: 10,
+             opset_import: [ "" : 19]
+           >
+           agraph (float[N] X) => (float[N] C)
+           <
+             float[1] weight = {1}
+           >
+           {
+              C = Cast<to=1>(X)
+           }
+        """
+        model = onnx.parser.parse_model(text_model)
+        model_repr = repr(model)
+        self.assertEqual(
+            model_repr,
+            "ModelProto(ir_version=10, opset_import={'': 19}, graph=GraphProto('agraph', input=<1 inputs>, output=<1 outputs>, initializer=<1 initializers>, node=<1 nodes>))",
+        )
+
+        graph_repr = repr(model.graph)
+        self.assertEqual(
+            graph_repr,
+            "GraphProto('agraph', input=<1 inputs>, output=<1 outputs>, initializer=<1 initializers>, node=<1 nodes>)",
+        )
+
+    def test_function_repr(self) -> None:
+        text = """
+            <
+            ir_version: 9,
+            opset_import: [ "" : 15, "custom_domain" : 1],
+            producer_name: "FunctionProtoTest",
+            producer_version: "1.0",
+            model_version: 1,
+            doc_string: "A test model for model local functions."
+          >
+         agraph (float[N] x) => (float[N] out)
+         {
+            out = custom_domain.Selu<alpha=2.0, gamma=3.0>(x)
+         }
+         <
+         domain: "custom_domain",
+         opset_import: [ "" : 15],
+         doc_string: "Test function proto"
+         >
+           Selu
+           <alpha: float=1.67326319217681884765625, gamma: float=1.05070102214813232421875>
+           (X) => (C)
+           {
+               constant_alpha = Constant<value_float: float=@alpha>()
+               constant_gamma = Constant<value_float: float=@gamma>()
+               alpha_x = CastLike(constant_alpha, X)
+               gamma_x = CastLike(constant_gamma, X)
+               exp_x = Exp(X)
+               alpha_x_exp_x = Mul(alpha_x, exp_x)
+               alpha_x_exp_x_ = Sub(alpha_x_exp_x, alpha_x)
+               neg = Mul(gamma_x, alpha_x_exp_x_)
+               pos = Mul(gamma_x, X)
+               _zero = Constant<value_float=0.0>()
+               zero = CastLike(_zero, X)
+               less_eq = LessOrEqual(X, zero)
+               C = Where(less_eq, neg, pos)
+           }
+        """
+        model = onnx.parser.parse_model(text)
+        self.assertEqual(
+            repr(model),
+            "ModelProto(ir_version=9, opset_import={'': 15, 'custom_domain': 1}, producer_name='FunctionProtoTest', producer_version='1.0', graph=GraphProto('agraph', input=<1 inputs>, output=<1 outputs>, node=<1 nodes>), functions=<1 functions>)",
+        )
+        function_repr = repr(model.functions[0])
+        self.assertEqual(
+            function_repr,
+            "FunctionProto('Selu', domain='custom_domain', opset_import={'': 15}, input=<1 inputs>, output=<1 outputs>, node=<13 nodes>)",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

onnx/test/compose_test.py

@@ -393,7 +393,96 @@ class TestComposeFunctions(unittest.TestCase):
         m3 = compose.merge_models(m1, m2, io_map=io_map)
         checker.check_model(m3)
 
-    # FIXME: This function should be removed, as tests should not contain a copy of the tested logic.
+    def test_add_prefix_to_inputs_outputs(self) -> None:
+        """Tests prefixing inputs and outputs nodes."""
+        input_graph = """
+            <
+                ir_version: 6,
+                opset_import: [ "": 13]
+            >
+            agraph (float[N, 128] X, float[128, 10] W, float[10] B) => (float[N, 10] C)
+            {
+                T = MatMul(X, W)
+                S = Add(T, B)
+                C = Softmax(S)
+            }
+            """
+        input_model = _load_model(input_graph)
+        prefix = "pre_"
+        prefixed_model = compose.add_prefix(
+            model=input_model,
+            prefix=prefix,
+            rename_inputs=True,
+            rename_outputs=True,
+            rename_edges=True,
+        )
+
+        checker.check_graph(prefixed_model.graph)
+
+        for i in prefixed_model.graph.input:
+            self.assertTrue(i.name.startswith(prefix))
+        for o in prefixed_model.graph.output:
+            self.assertTrue(o.name.startswith(prefix))
+
+    def test_add_prefix_wo_inputs_outputs(self) -> None:
+        """Tests prefixing input and output nodes, when renaming of inputs/outputs is deactivated."""
+        input_graph = """
+            <
+                ir_version: 6,
+                opset_import: [ "": 13]
+            >
+            agraph (float[2,2] A) => (float[2,2] B)
+            {
+            B = Identity(A)
+            }
+            """
+        input_model = _load_model(input_graph)
+        prefix = "pre_"
+        prefixed_model = compose.add_prefix(
+            model=input_model,
+            prefix=prefix,
+            rename_inputs=False,
+            rename_outputs=False,
+            rename_edges=True,
+        )
+
+        checker.check_graph(prefixed_model.graph)
+        self.assertEqual(input_model.graph.input, prefixed_model.graph.input)
+        self.assertEqual(input_model.graph.output, prefixed_model.graph.output)
+
+    def test_add_prefix_with_loose_inputs_outputs(self) -> None:
+        """Tests prefixing of graphs with loose inputs."""
+        input_graph = """
+            <
+                ir_version: 7,
+                opset_import: [ "" : 13 ]
+            >
+            agraph (bool b, float[128] X, float[128] Y, bool unused_input) => (float[128] Z)
+            {
+            Z = If (b) <
+                then_branch = g1 () => (float[128] z_then) { z_then = Add(X, X) },
+                else_branch = g2 () => (float[128] z_else) { z_else = Sub(Y, Y) }
+                >
+            }
+            """
+        input_model = _load_model(input_graph)
+        prefix = "pre_"
+        prefixed_model = compose.add_prefix(
+            model=input_model,
+            prefix=prefix,
+            rename_inputs=True,
+            rename_outputs=True,
+            rename_edges=True,
+        )
+        checker.check_graph(prefixed_model.graph)
+
+        # Check if all inputs are prefixed even if input is not connected to node. e.g., 'unused_input'.
+        for i in prefixed_model.graph.input:
+            self.assertTrue(i.name.startswith(prefix))
+
+        for graph_output in prefixed_model.graph.output:
+            self.assertTrue(graph_output.name.startswith(prefix))
+
     def _test_add_prefix(
         self,
         rename_nodes: bool = False,
@@ -516,7 +605,9 @@ class TestComposeFunctions(unittest.TestCase):
 
     def test_add_prefix_edges(self) -> None:
         """Tests prefixing nodes edges. This will also rename inputs/outputs, since the names are shared"""
-        self._test_add_prefix(rename_edges=True)
+        self._test_add_prefix(
+            rename_edges=True, rename_inputs=True, rename_outputs=True
+        )
 
     def test_add_prefix_inputs(self) -> None:
         """Tests prefixing graph inputs only. Relevant node edges should be renamed as well"""

onnx/test/cpp/function_context_test.cc

@@ -71,9 +71,10 @@ BuildFloatFunctionBody(const FunctionBodyBuildContext& /*ctx*/, const OpSchema&
   // Create a scalar-tensor constant 2.0 of float type:
   auto two_as_tensor = ToTensor(2.0, TensorProto_DataType::TensorProto_DataType_FLOAT);
 
-  std::vector<FunctionBodyHelper::NodeDef> body{// nodes: {outputs, op, inputs, attributes}
-                                                {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
-                                                {{"Y"}, "Mul", {"X", "Two"}}};
+  std::vector<FunctionBodyHelper::NodeDef> body{
+      // nodes: {outputs, op, inputs, attributes}
+      {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
+      {{"Y"}, "Mul", {"X", "Two"}}};
 
   return BuildFunctionProto(functionProto, schema, body);
 }
@@ -92,7 +93,7 @@ static void RegisterCustomFuncFloatSchema() {
   (void)unused;
 }
 
-// Test for Context dependant function without type context
+// Test for Context dependent function without type context
 TEST(FunctionAPITest, ContextDependentFunctionTest) {
   RegisterCustomFuncFloatSchema();
 
@@ -130,9 +131,10 @@ BuildFunctionBody(const FunctionBodyBuildContext& ctx, const OpSchema& schema, F
   auto elem_type = (TensorProto_DataType)tp->tensor_type().elem_type();
   auto two_as_tensor = ToTensor(2.0, elem_type);
 
-  std::vector<FunctionBodyHelper::NodeDef> body{// nodes: {outputs, op, inputs, attributes}
-                                                {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
-                                                {{"Y"}, "Mul", {"X", "Two"}}};
+  std::vector<FunctionBodyHelper::NodeDef> body{
+      // nodes: {outputs, op, inputs, attributes}
+      {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
+      {{"Y"}, "Mul", {"X", "Two"}}};
 
   return BuildFunctionProto(functionProto, schema, body);
 }