onnx 1.20.1__cp311-cp311-win_arm64.whl

onnx/serialization.py

@@ -0,0 +1,212 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import warnings
+
+__all__ = [
+    "registry",
+]
+
+import typing
+from typing import Any, Protocol, TypeVar
+
+import google.protobuf.json_format
+import google.protobuf.message
+import google.protobuf.text_format
+
+import onnx
+
+if typing.TYPE_CHECKING:
+    from collections.abc import Collection
+
+_Proto = TypeVar("_Proto", bound=google.protobuf.message.Message)
+# Encoding used for serializing and deserializing text files
+_ENCODING = "utf-8"
+
+
+class ProtoSerializer(Protocol):
+    """A serializer-deserializer to and from in-memory Protocol Buffers representations."""
+
+    # Format supported by the serializer. E.g. "protobuf"
+    supported_format: str
+    # File extensions supported by the serializer. E.g. frozenset({".onnx", ".pb"})
+    # Be careful to include the dot in the file extension.
+    file_extensions: Collection[str]
+
+    # NOTE: The methods defined are serialize_proto and deserialize_proto and not the
+    # more generic serialize and deserialize to leave space for future protocols
+    # that are defined to serialize/deserialize the ONNX in memory IR.
+    # This way a class can implement both protocols.
+
+    def serialize_proto(self, proto: _Proto) -> Any:
+        """Serialize a in-memory proto to a serialized data type."""
+
+    def deserialize_proto(self, serialized: Any, proto: _Proto) -> _Proto:
+        """Parse a serialized data type into a in-memory proto."""
+
+
+class _Registry:
+    def __init__(self) -> None:
+        self._serializers: dict[str, ProtoSerializer] = {}
+        # A mapping from file extension to format
+        self._extension_to_format: dict[str, str] = {}
+
+    def register(self, serializer: ProtoSerializer) -> None:
+        self._serializers[serializer.supported_format] = serializer
+        self._extension_to_format.update(
+            dict.fromkeys(serializer.file_extensions, serializer.supported_format)
+        )
+
+    def get(self, fmt: str) -> ProtoSerializer:
+        """Get a serializer for a format.
+
+        Args:
+            fmt: The format to get a serializer for.
+
+        Returns:
+            ProtoSerializer: The serializer for the format.
+
+        Raises:
+            ValueError: If the format is not supported.
+        """
+        try:
+            return self._serializers[fmt]
+        except KeyError:
+            raise ValueError(
+                f"Unsupported format: '{fmt}'. Supported formats are: {self._serializers.keys()}"
+            ) from None
+
+    def get_format_from_file_extension(self, file_extension: str) -> str | None:
+        """Get the corresponding format from a file extension.
+
+        Args:
+            file_extension: The file extension to get a format for.
+
+        Returns:
+            The format for the file extension, or None if not found.
+        """
+        return self._extension_to_format.get(file_extension)
+
+
+class _ProtobufSerializer(ProtoSerializer):
+    """Serialize and deserialize protobuf message."""
+
+    supported_format = "protobuf"
+    file_extensions = frozenset({".onnx", ".pb"})
+
+    def serialize_proto(self, proto: _Proto) -> bytes:
+        if hasattr(proto, "SerializeToString") and callable(proto.SerializeToString):
+            try:
+                result = proto.SerializeToString()
+            except ValueError as e:
+                if proto.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF:
+                    raise ValueError(
+                        "The proto size is larger than the 2 GB limit. "
+                        "Please use save_as_external_data to save tensors separately from the model file."
+                    ) from e
+                raise
+            return result  # type: ignore
+        raise TypeError(
+            f"No SerializeToString method is detected.\ntype is {type(proto)}"
+        )
+
+    def deserialize_proto(self, serialized: bytes, proto: _Proto) -> _Proto:
+        if not isinstance(serialized, bytes):
+            raise TypeError(
+                f"Parameter 'serialized' must be bytes, but got type: {type(serialized)}"
+            )
+        decoded = typing.cast("int | None", 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)}"
+            )
+        return proto
+
+
+class _TextProtoSerializer(ProtoSerializer):
+    """Serialize and deserialize text proto."""
+
+    supported_format = "textproto"
+    file_extensions = frozenset({".txtpb", ".textproto", ".prototxt", ".pbtxt"})
+
+    def serialize_proto(self, proto: _Proto) -> bytes:
+        textproto = google.protobuf.text_format.MessageToString(proto)
+        return textproto.encode(_ENCODING)
+
+    def deserialize_proto(self, serialized: bytes | str, proto: _Proto) -> _Proto:
+        if not isinstance(serialized, (bytes, str)):
+            raise TypeError(
+                f"Parameter 'serialized' must be bytes or str, but got type: {type(serialized)}"
+            )
+        if isinstance(serialized, bytes):
+            serialized = serialized.decode(_ENCODING)
+        assert isinstance(serialized, str)
+        return google.protobuf.text_format.Parse(serialized, proto)
+
+
+class _JsonSerializer(ProtoSerializer):
+    """Serialize and deserialize JSON."""
+
+    supported_format = "json"
+    file_extensions = frozenset({".json", ".onnxjson"})
+
+    def serialize_proto(self, proto: _Proto) -> bytes:
+        json_message = google.protobuf.json_format.MessageToJson(
+            proto, preserving_proto_field_name=True
+        )
+        return json_message.encode(_ENCODING)
+
+    def deserialize_proto(self, serialized: bytes | str, proto: _Proto) -> _Proto:
+        if not isinstance(serialized, (bytes, str)):
+            raise TypeError(
+                f"Parameter 'serialized' must be bytes or str, but got type: {type(serialized)}"
+            )
+        if isinstance(serialized, bytes):
+            serialized = serialized.decode(_ENCODING)
+        assert isinstance(serialized, str)
+        return google.protobuf.json_format.Parse(serialized, proto)
+
+
+class _TextualSerializer(ProtoSerializer):
+    """Serialize and deserialize the ONNX textual representation."""
+
+    supported_format = "onnxtxt"
+    file_extensions = frozenset({".onnxtxt", ".onnxtext"})
+
+    def serialize_proto(self, proto: _Proto) -> bytes:
+        text = onnx.printer.to_text(proto)  # type: ignore[arg-type]
+        return text.encode(_ENCODING)
+
+    def deserialize_proto(self, serialized: bytes | str, proto: _Proto) -> _Proto:
+        warnings.warn(
+            "The onnxtxt format is experimental. Please report any errors to the ONNX GitHub repository.",
+            stacklevel=2,
+        )
+        if not isinstance(serialized, (bytes, str)):
+            raise TypeError(
+                f"Parameter 'serialized' must be bytes or str, but got type: {type(serialized)}"
+            )
+        if isinstance(serialized, bytes):
+            text = serialized.decode(_ENCODING)
+        else:
+            text = serialized
+        if isinstance(proto, onnx.ModelProto):
+            return onnx.parser.parse_model(text)  # type: ignore[return-value]
+        if isinstance(proto, onnx.GraphProto):
+            return onnx.parser.parse_graph(text)  # type: ignore[return-value]
+        if isinstance(proto, onnx.FunctionProto):
+            return onnx.parser.parse_function(text)  # type: ignore[return-value]
+        if isinstance(proto, onnx.NodeProto):
+            return onnx.parser.parse_node(text)  # type: ignore[return-value]
+        raise ValueError(f"Unsupported proto type: {type(proto)}")
+
+
+# Register default serializers
+registry = _Registry()
+registry.register(_ProtobufSerializer())
+registry.register(_TextProtoSerializer())
+registry.register(_JsonSerializer())
+registry.register(_TextualSerializer())

onnx/shape_inference/attribute_binder.h

@@ -0,0 +1,67 @@
+// Copyright (c) ONNX Project Contributors
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+#include "onnx/common/visitor.h"
+
+namespace ONNX_NAMESPACE {
+namespace internal { // internal/private API
+
+using AttributeMap = std::unordered_map<std::string, const AttributeProto*>;
+
+// Class for binding formal attribute-parameters (in a node or graph) to their values.
+
+class AttributeBinder : public MutableVisitor {
+ public:
+  explicit AttributeBinder(const AttributeMap& attr_map) : attr_map_(attr_map) {}
+
+  // Binding a formal attribute-parameter to a value may, as a special case, also
+  // remove the attribute from the list of attributes of a node (when the attribute
+  // has no specified value). Hence, we need to do the processing at a Node level
+  // rather than an attribute level.
+  void VisitNode(NodeProto* node) override {
+    auto& attributes = *node->mutable_attribute();
+    for (auto attr_iter = attributes.begin(); attr_iter != attributes.end();) {
+      auto& attr = *attr_iter;
+      if (!attr.ref_attr_name().empty()) {
+        // Attribute-references must be replaced by the corresponding attribute-value in the call-node
+        // if the call-node contains the attribute. Otherwise, this attribute must be removed.
+        auto it = attr_map_.find(attr.ref_attr_name());
+        if (it != attr_map_.end()) {
+          const AttributeProto* replacement = it->second;
+          // Copy value of attribute, but retain original name:
+          std::string name = attr.name();
+          attr = *replacement;
+          attr.set_name(name);
+          ++attr_iter;
+        } else {
+          attr_iter = attributes.erase(attr_iter);
+        }
+      } else {
+        // For regular attributes, we process subgraphs, if present, recursively.
+        VisitAttribute(&attr);
+        ++attr_iter;
+      }
+    }
+  }
+
+  // Updates a FunctionProto by replacing all attribute-references with the corresponding
+  // attribute-values in the call-node, if present. Otherwise, the attribute is removed.
+  static void BindAttributes(const NodeProto& callnode, FunctionProto& callee) {
+    AttributeMap map;
+    for (auto& attr : callnode.attribute()) {
+      map[attr.name()] = &attr;
+    }
+    AttributeBinder attr_binder(map);
+    attr_binder.VisitFunction(&callee);
+  }
+
+ private:
+  const AttributeMap& attr_map_;
+};
+
+} // namespace internal
+} // namespace ONNX_NAMESPACE