onnx 1.17.0__cp311-cp311-macosx_12_0_universal2.whl → 1.18.0__cp311-cp311-macosx_12_0_universal2.whl

onnx/onnx.in.proto

@@ -107,9 +107,13 @@ enum Version {
   // Added FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ.
   IR_VERSION_2023_5_5 = 0x0000000000000009;
 
-  // IR VERSION 10 published on TBD
+  // IR VERSION 10 published on March 25, 2024
   // Added UINT4, INT4.
-  IR_VERSION = 0x000000000000000A;
+  IR_VERSION_2024_3_25 = 0x000000000000000A;
+
+  // IR VERSION 11 published on TBD
+  // Added FLOAT4E2M1, multi-device protobuf classes.
+  IR_VERSION = 0x000000000000000B;
 }
 
 // Attributes
@@ -227,6 +231,83 @@ message NodeProto {
 
   // Named metadata values; keys should be distinct.
   repeated StringStringEntryProto metadata_props = 9;
+
+  // Configuration of multi-device annotations.
+  repeated NodeDeviceConfigurationProto device_configurations = 10;
+}
+
+// IntIntListEntryProto follows the pattern for cross-proto-version maps.
+// See https://developers.google.com/protocol-buffers/docs/proto3#maps
+message IntIntListEntryProto {
+  optional int64 key = 1;
+  repeated int64 value = 2;
+};
+
+// Multi-device configuration proto for NodeProto.
+message NodeDeviceConfigurationProto {
+    // This field MUST be present for this version of the IR.
+    // ID of the configuration. MUST match the name of a DeviceConfigurationProto.
+    optional string configuration_id = 1;
+    // Sharding spec for the node.
+    repeated ShardingSpecProto sharding_spec = 2;
+    // Pipeline stage of this node.
+    optional int32 pipeline_stage = 3;
+}
+
+// ShardingSpecProto: This describes the sharding spec for a specific
+// input or output tensor of a node.
+message ShardingSpecProto {
+  // This field MUST be present for this version of the IR.
+  // Identifies the input or output of the node that is being sharded.
+  // Required to match a name specified in the node's input or output list of ValueInfoProtos.
+  // It is called `logical tensor` in subsequent descriptions.
+  optional string tensor_name = 1;
+
+  // The following is the list of devices across which the logical
+  // tensor is sharded or replicated.
+  repeated int64 device = 2;
+
+  // Each element v in above field devices may represent either a
+  // device or a set of devices (when we want the same shard/tensor
+  // to be replicated across a subset of devices), as indicated by
+  // the following optional map. If the map contains an entry for v,
+  // then v represents a device group, and the map indicates the set
+  // of devices in that group.
+  repeated IntIntListEntryProto index_to_device_group_map = 3;
+
+  // The following is the sharded-shape of the tensor, consisting of
+  // the sharding-spec for each axis of the tensor.
+  repeated ShardedDimProto sharded_dim = 4;
+}
+
+// ShardedDimProto: This describes the sharding spec for a single
+// axis of a sharded tensor.
+message ShardedDimProto {
+  // This field MUST be present for this version of the IR.
+  // The axis this sharding corresponds to. Must be in the range of
+  // [-r, r - 1], where r is the rank of the tensor. Negative axis values means
+  // counting from the back.
+  optional int64 axis = 1;
+
+  // Describes how the tensor on the provided axis is sharded.
+  // The common-case is described by a single instance of SimpleShardedDimProto.
+  // Multiple instances can be used to handle cases where a sharded
+  // tensor is reshaped, fusing multiple axes into one.
+  repeated SimpleShardedDimProto simple_sharding = 2;
+}
+
+// SimpleShardedDimProto: Indicates that N blocks are divided into M shards.
+// N is allowed to be symbolic where M is required to be a constant.
+message SimpleShardedDimProto {
+    // Dimension value to be sharded.
+    oneof dim {
+        int64 dim_value = 1;
+        string dim_param = 2;
+    }
+
+    // This field MUST be present for this version of the IR.
+    // Number of shards to split dim into.
+    optional int64 num_shards = 3;
 }
 
 // Training information
@@ -426,8 +507,24 @@ message ModelProto {
   // One FunctionProto can reference other FunctionProto in the model, however, recursive reference
   // is not allowed.
   repeated FunctionProto functions = 25;
+
+  // Describes different target configurations for a multi-device use case.
+  // A model MAY describe multiple multi-device configurations for execution.
+  repeated DeviceConfigurationProto configuration = 26;
 };
 
+// DeviceConfigurationProto describes a multi-device configuration for a model.
+message DeviceConfigurationProto {
+    // This field MUST be present for this version of the IR.
+    // Name of the configuration.
+    optional string name = 1;
+    // This field MUST be present for this version of the IR.
+    // Number of devices inside this configuration.
+    optional int32 num_devices = 2;
+    // Optional names of the devices. MUST be length of num_devices if provided.
+    repeated string device = 3;
+}
+
 // StringStringEntryProto follows the pattern for cross-proto-version maps.
 // See https://developers.google.com/protocol-buffers/docs/proto3#maps
 message StringStringEntryProto {
@@ -535,10 +632,13 @@ message TensorProto {
     FLOAT8E5M2 = 19;      // follows IEEE 754, supports nan, inf, mostly used for gradients
     FLOAT8E5M2FNUZ = 20;  // follows IEEE 754, supports nan, not inf, mostly used for gradients, no negative zero
 
-    // 4-bit data-types
+    // 4-bit integer data types
     UINT4 = 21;  // Unsigned integer in range [0, 15]
     INT4 = 22;   // Signed integer in range [-8, 7], using two's-complement representation
 
+    // 4-bit floating point data types
+    FLOAT4E2M1 = 23;
+
     // Future extensions go here.
   }
 
@@ -572,13 +672,19 @@ message TensorProto {
   // When this field is present, the data_type field MUST be FLOAT or COMPLEX64.
   repeated float float_data = 4 [packed = true];
 
-  // For int32, uint8, int8, uint16, int16, uint4, int4, bool, float8 and float16 values
-  // float16 and float8 values must be bit-wise converted to an uint16_t prior
-  // to writing to the buffer.
-  // uint4 and int4 values must be packed to 4bitx2 prior to writing to the buffer, the first element is stored in
-  // the 4 LSB and the second element is stored in the 4 MSB.
+  // For int32, uint8, int8, uint16, int16, uint4, int4, bool, (b)float16, float8, and float4:
+  // - (b)float16 and float8 values MUST be converted bit-wise into an unsigned integer
+  //   representation before being written to the buffer.
+  // - Each pair of uint4, int4, and float4 values MUST be packed as two 4-bit elements into a single byte.
+  //   The first element is stored in the 4 least significant bits (LSB),
+  //   and the second element is stored in the 4 most significant bits (MSB).
+  //
+  // Consequently:
+  // - For data types with a bit-width of 8 or greater, each `int32_data` stores one element.
+  // - For 4-bit data types, each `int32_data` stores two elements.
+  //
   // When this field is present, the data_type field MUST be
-  // INT32, INT16, INT8, INT4, UINT16, UINT8, UINT4, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ
+  // INT32, INT16, INT8, INT4, UINT16, UINT8, UINT4, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ, FLOAT4E2M1
   repeated int32 int32_data = 5 [packed = true];
 
   // For strings.

onnx/onnx.proto

@@ -108,9 +108,13 @@ enum Version {
   // Added FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ.
   IR_VERSION_2023_5_5 = 0x0000000000000009;
 
-  // IR VERSION 10 published on TBD
+  // IR VERSION 10 published on March 25, 2024
   // Added UINT4, INT4.
-  IR_VERSION = 0x000000000000000A;
+  IR_VERSION_2024_3_25 = 0x000000000000000A;
+
+  // IR VERSION 11 published on TBD
+  // Added FLOAT4E2M1, multi-device protobuf classes.
+  IR_VERSION = 0x000000000000000B;
 }
 
 // Attributes
@@ -228,6 +232,83 @@ message NodeProto {
 
   // Named metadata values; keys should be distinct.
   repeated StringStringEntryProto metadata_props = 9;
+
+  // Configuration of multi-device annotations.
+  repeated NodeDeviceConfigurationProto device_configurations = 10;
+}
+
+// IntIntListEntryProto follows the pattern for cross-proto-version maps.
+// See https://developers.google.com/protocol-buffers/docs/proto3#maps
+message IntIntListEntryProto {
+  optional int64 key = 1;
+  repeated int64 value = 2;
+};
+
+// Multi-device configuration proto for NodeProto.
+message NodeDeviceConfigurationProto {
+    // This field MUST be present for this version of the IR.
+    // ID of the configuration. MUST match the name of a DeviceConfigurationProto.
+    optional string configuration_id = 1;
+    // Sharding spec for the node.
+    repeated ShardingSpecProto sharding_spec = 2;
+    // Pipeline stage of this node.
+    optional int32 pipeline_stage = 3;
+}
+
+// ShardingSpecProto: This describes the sharding spec for a specific
+// input or output tensor of a node.
+message ShardingSpecProto {
+  // This field MUST be present for this version of the IR.
+  // Identifies the input or output of the node that is being sharded.
+  // Required to match a name specified in the node's input or output list of ValueInfoProtos.
+  // It is called `logical tensor` in subsequent descriptions.
+  optional string tensor_name = 1;
+
+  // The following is the list of devices across which the logical
+  // tensor is sharded or replicated.
+  repeated int64 device = 2;
+
+  // Each element v in above field devices may represent either a
+  // device or a set of devices (when we want the same shard/tensor
+  // to be replicated across a subset of devices), as indicated by
+  // the following optional map. If the map contains an entry for v,
+  // then v represents a device group, and the map indicates the set
+  // of devices in that group.
+  repeated IntIntListEntryProto index_to_device_group_map = 3;
+
+  // The following is the sharded-shape of the tensor, consisting of
+  // the sharding-spec for each axis of the tensor.
+  repeated ShardedDimProto sharded_dim = 4;
+}
+
+// ShardedDimProto: This describes the sharding spec for a single
+// axis of a sharded tensor.
+message ShardedDimProto {
+  // This field MUST be present for this version of the IR.
+  // The axis this sharding corresponds to. Must be in the range of
+  // [-r, r - 1], where r is the rank of the tensor. Negative axis values means
+  // counting from the back.
+  optional int64 axis = 1;
+
+  // Describes how the tensor on the provided axis is sharded.
+  // The common-case is described by a single instance of SimpleShardedDimProto.
+  // Multiple instances can be used to handle cases where a sharded
+  // tensor is reshaped, fusing multiple axes into one.
+  repeated SimpleShardedDimProto simple_sharding = 2;
+}
+
+// SimpleShardedDimProto: Indicates that N blocks are divided into M shards.
+// N is allowed to be symbolic where M is required to be a constant.
+message SimpleShardedDimProto {
+    // Dimension value to be sharded.
+    oneof dim {
+        int64 dim_value = 1;
+        string dim_param = 2;
+    }
+
+    // This field MUST be present for this version of the IR.
+    // Number of shards to split dim into.
+    optional int64 num_shards = 3;
 }
 
 // Training information
@@ -427,8 +508,24 @@ message ModelProto {
   // One FunctionProto can reference other FunctionProto in the model, however, recursive reference
   // is not allowed.
   repeated FunctionProto functions = 25;
+
+  // Describes different target configurations for a multi-device use case.
+  // A model MAY describe multiple multi-device configurations for execution.
+  repeated DeviceConfigurationProto configuration = 26;
 };
 
+// DeviceConfigurationProto describes a multi-device configuration for a model.
+message DeviceConfigurationProto {
+    // This field MUST be present for this version of the IR.
+    // Name of the configuration.
+    optional string name = 1;
+    // This field MUST be present for this version of the IR.
+    // Number of devices inside this configuration.
+    optional int32 num_devices = 2;
+    // Optional names of the devices. MUST be length of num_devices if provided.
+    repeated string device = 3;
+}
+
 // StringStringEntryProto follows the pattern for cross-proto-version maps.
 // See https://developers.google.com/protocol-buffers/docs/proto3#maps
 message StringStringEntryProto {
@@ -536,10 +633,13 @@ message TensorProto {
     FLOAT8E5M2 = 19;      // follows IEEE 754, supports nan, inf, mostly used for gradients
     FLOAT8E5M2FNUZ = 20;  // follows IEEE 754, supports nan, not inf, mostly used for gradients, no negative zero
 
-    // 4-bit data-types
+    // 4-bit integer data types
     UINT4 = 21;  // Unsigned integer in range [0, 15]
     INT4 = 22;   // Signed integer in range [-8, 7], using two's-complement representation
 
+    // 4-bit floating point data types
+    FLOAT4E2M1 = 23;
+
     // Future extensions go here.
   }
 
@@ -573,13 +673,19 @@ message TensorProto {
   // When this field is present, the data_type field MUST be FLOAT or COMPLEX64.
   repeated float float_data = 4 [packed = true];
 
-  // For int32, uint8, int8, uint16, int16, uint4, int4, bool, float8 and float16 values
-  // float16 and float8 values must be bit-wise converted to an uint16_t prior
-  // to writing to the buffer.
-  // uint4 and int4 values must be packed to 4bitx2 prior to writing to the buffer, the first element is stored in
-  // the 4 LSB and the second element is stored in the 4 MSB.
+  // For int32, uint8, int8, uint16, int16, uint4, int4, bool, (b)float16, float8, and float4:
+  // - (b)float16 and float8 values MUST be converted bit-wise into an unsigned integer
+  //   representation before being written to the buffer.
+  // - Each pair of uint4, int4, and float4 values MUST be packed as two 4-bit elements into a single byte.
+  //   The first element is stored in the 4 least significant bits (LSB),
+  //   and the second element is stored in the 4 most significant bits (MSB).
+  //
+  // Consequently:
+  // - For data types with a bit-width of 8 or greater, each `int32_data` stores one element.
+  // - For 4-bit data types, each `int32_data` stores two elements.
+  //
   // When this field is present, the data_type field MUST be
-  // INT32, INT16, INT8, INT4, UINT16, UINT8, UINT4, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ
+  // INT32, INT16, INT8, INT4, UINT16, UINT8, UINT4, BOOL, FLOAT16, BFLOAT16, FLOAT8E4M3FN, FLOAT8E4M3FNUZ, FLOAT8E5M2, FLOAT8E5M2FNUZ, FLOAT4E2M1
   repeated int32 int32_data = 5 [packed = true];
 
   // For strings.

onnx/onnx_cpp2py_export/defs.pyi

@@ -1,5 +1,6 @@
 """Submodule containing all the ONNX schema definitions."""
-from typing import Sequence, overload
+from collections.abc import Sequence
+from typing import overload
 
 from onnx import AttributeProto, FunctionProto
 
@@ -148,7 +149,7 @@ class OpSchema:
         def __init__(
             self,
             name: str,
-            type: OpSchema.AttrType,  # noqa: A002
+            type: OpSchema.AttrType,
             description: str = "",
             *,
             required: bool = True,

onnx/onnx_data_pb2.py

@@ -1,11 +1,12 @@
 # -*- coding: utf-8 -*-
 # Generated by the protocol buffer compiler.  DO NOT EDIT!
 # source: onnx/onnx-data.proto
+# Protobuf Python Version: 4.25.1
 """Generated protocol buffer code."""
-from google.protobuf.internal import builder as _builder
 from google.protobuf import descriptor as _descriptor
 from google.protobuf import descriptor_pool as _descriptor_pool
 from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
 # @@protoc_insertion_point(imports)
 
 _sym_db = _symbol_database.Default()
@@ -16,20 +17,20 @@ from onnx import onnx_ml_pb2 as onnx_dot_onnx__ml__pb2
 
 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x14onnx/onnx-data.proto\x12\x04onnx\x1a\x12onnx/onnx-ml.proto\"\xf0\x02\n\rSequenceProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x11\n\telem_type\x18\x02 \x01(\x05\x12(\n\rtensor_values\x18\x03 \x03(\x0b\x32\x11.onnx.TensorProto\x12\x35\n\x14sparse_tensor_values\x18\x04 \x03(\x0b\x32\x17.onnx.SparseTensorProto\x12,\n\x0fsequence_values\x18\x05 \x03(\x0b\x32\x13.onnx.SequenceProto\x12\"\n\nmap_values\x18\x06 \x03(\x0b\x32\x0e.onnx.MapProto\x12,\n\x0foptional_values\x18\x07 \x03(\x0b\x32\x13.onnx.OptionalProto\"]\n\x08\x44\x61taType\x12\r\n\tUNDEFINED\x10\x00\x12\n\n\x06TENSOR\x10\x01\x12\x11\n\rSPARSE_TENSOR\x10\x02\x12\x0c\n\x08SEQUENCE\x10\x03\x12\x07\n\x03MAP\x10\x04\x12\x0c\n\x08OPTIONAL\x10\x05\"r\n\x08MapProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x10\n\x08key_type\x18\x02 \x01(\x05\x12\x0c\n\x04keys\x18\x03 \x03(\x03\x12\x13\n\x0bstring_keys\x18\x04 \x03(\x0c\x12#\n\x06values\x18\x05 \x01(\x0b\x32\x13.onnx.SequenceProto\"\xeb\x02\n\rOptionalProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x11\n\telem_type\x18\x02 \x01(\x05\x12\'\n\x0ctensor_value\x18\x03 \x01(\x0b\x32\x11.onnx.TensorProto\x12\x34\n\x13sparse_tensor_value\x18\x04 \x01(\x0b\x32\x17.onnx.SparseTensorProto\x12+\n\x0esequence_value\x18\x05 \x01(\x0b\x32\x13.onnx.SequenceProto\x12!\n\tmap_value\x18\x06 \x01(\x0b\x32\x0e.onnx.MapProto\x12+\n\x0eoptional_value\x18\x07 \x01(\x0b\x32\x13.onnx.OptionalProto\"]\n\x08\x44\x61taType\x12\r\n\tUNDEFINED\x10\x00\x12\n\n\x06TENSOR\x10\x01\x12\x11\n\rSPARSE_TENSOR\x10\x02\x12\x0c\n\x08SEQUENCE\x10\x03\x12\x07\n\x03MAP\x10\x04\x12\x0c\n\x08OPTIONAL\x10\x05\x42\x02H\x03')
 
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
-_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'onnx.onnx_data_pb2', globals())
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'onnx.onnx_data_pb2', _globals)
 if _descriptor._USE_C_DESCRIPTORS == False:
-
-  DESCRIPTOR._options = None
-  DESCRIPTOR._serialized_options = b'H\003'
-  _SEQUENCEPROTO._serialized_start=51
-  _SEQUENCEPROTO._serialized_end=419
-  _SEQUENCEPROTO_DATATYPE._serialized_start=326
-  _SEQUENCEPROTO_DATATYPE._serialized_end=419
-  _MAPPROTO._serialized_start=421
-  _MAPPROTO._serialized_end=535
-  _OPTIONALPROTO._serialized_start=538
-  _OPTIONALPROTO._serialized_end=901
-  _OPTIONALPROTO_DATATYPE._serialized_start=326
-  _OPTIONALPROTO_DATATYPE._serialized_end=419
+  _globals['DESCRIPTOR']._options = None
+  _globals['DESCRIPTOR']._serialized_options = b'H\003'
+  _globals['_SEQUENCEPROTO']._serialized_start=51
+  _globals['_SEQUENCEPROTO']._serialized_end=419
+  _globals['_SEQUENCEPROTO_DATATYPE']._serialized_start=326
+  _globals['_SEQUENCEPROTO_DATATYPE']._serialized_end=419
+  _globals['_MAPPROTO']._serialized_start=421
+  _globals['_MAPPROTO']._serialized_end=535
+  _globals['_OPTIONALPROTO']._serialized_start=538
+  _globals['_OPTIONALPROTO']._serialized_end=901
+  _globals['_OPTIONALPROTO_DATATYPE']._serialized_start=326
+  _globals['_OPTIONALPROTO_DATATYPE']._serialized_end=419
 # @@protoc_insertion_point(module_scope)

onnx/onnx_data_pb2.pyi

@@ -1,146 +1,86 @@
-# @generated by protoc-gen-mypy.py.  Do not edit!
-# mypy: disable-error-code=override
-from google.protobuf.message import (  # type: ignore
-    Message,
-)
+from onnx import onnx_ml_pb2 as _onnx_ml_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
 
-from typing import (
-    Iterable,
-    List,
-    Optional as OptionalType,
-    Tuple,
-    cast,
-)
+DESCRIPTOR: _descriptor.FileDescriptor
 
-from google.protobuf.internal.containers import (  # type: ignore
-    RepeatedCompositeFieldContainer,
-    RepeatedScalarFieldContainer,
-)
+class SequenceProto(_message.Message):
+    __slots__ = ("name", "elem_type", "tensor_values", "sparse_tensor_values", "sequence_values", "map_values", "optional_values")
+    class DataType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = ()
+        UNDEFINED: _ClassVar[SequenceProto.DataType]
+        TENSOR: _ClassVar[SequenceProto.DataType]
+        SPARSE_TENSOR: _ClassVar[SequenceProto.DataType]
+        SEQUENCE: _ClassVar[SequenceProto.DataType]
+        MAP: _ClassVar[SequenceProto.DataType]
+        OPTIONAL: _ClassVar[SequenceProto.DataType]
+    UNDEFINED: SequenceProto.DataType
+    TENSOR: SequenceProto.DataType
+    SPARSE_TENSOR: SequenceProto.DataType
+    SEQUENCE: SequenceProto.DataType
+    MAP: SequenceProto.DataType
+    OPTIONAL: SequenceProto.DataType
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    ELEM_TYPE_FIELD_NUMBER: _ClassVar[int]
+    TENSOR_VALUES_FIELD_NUMBER: _ClassVar[int]
+    SPARSE_TENSOR_VALUES_FIELD_NUMBER: _ClassVar[int]
+    SEQUENCE_VALUES_FIELD_NUMBER: _ClassVar[int]
+    MAP_VALUES_FIELD_NUMBER: _ClassVar[int]
+    OPTIONAL_VALUES_FIELD_NUMBER: _ClassVar[int]
+    name: str
+    elem_type: int
+    tensor_values: _containers.RepeatedCompositeFieldContainer[_onnx_ml_pb2.TensorProto]
+    sparse_tensor_values: _containers.RepeatedCompositeFieldContainer[_onnx_ml_pb2.SparseTensorProto]
+    sequence_values: _containers.RepeatedCompositeFieldContainer[SequenceProto]
+    map_values: _containers.RepeatedCompositeFieldContainer[MapProto]
+    optional_values: _containers.RepeatedCompositeFieldContainer[OptionalProto]
+    def __init__(self, name: _Optional[str] = ..., elem_type: _Optional[int] = ..., tensor_values: _Optional[_Iterable[_Union[_onnx_ml_pb2.TensorProto, _Mapping]]] = ..., sparse_tensor_values: _Optional[_Iterable[_Union[_onnx_ml_pb2.SparseTensorProto, _Mapping]]] = ..., sequence_values: _Optional[_Iterable[_Union[SequenceProto, _Mapping]]] = ..., map_values: _Optional[_Iterable[_Union[MapProto, _Mapping]]] = ..., optional_values: _Optional[_Iterable[_Union[OptionalProto, _Mapping]]] = ...) -> None: ...
 
-from onnx.onnx_ml_pb2 import (
-    SparseTensorProto,
-    TensorProto,
-)
+class MapProto(_message.Message):
+    __slots__ = ("name", "key_type", "keys", "string_keys", "values")
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    KEY_TYPE_FIELD_NUMBER: _ClassVar[int]
+    KEYS_FIELD_NUMBER: _ClassVar[int]
+    STRING_KEYS_FIELD_NUMBER: _ClassVar[int]
+    VALUES_FIELD_NUMBER: _ClassVar[int]
+    name: str
+    key_type: int
+    keys: _containers.RepeatedScalarFieldContainer[int]
+    string_keys: _containers.RepeatedScalarFieldContainer[bytes]
+    values: SequenceProto
+    def __init__(self, name: _Optional[str] = ..., key_type: _Optional[int] = ..., keys: _Optional[_Iterable[int]] = ..., string_keys: _Optional[_Iterable[bytes]] = ..., values: _Optional[_Union[SequenceProto, _Mapping]] = ...) -> None: ...
 
-class SequenceProto(Message):
-    class DataType(int):
-        @classmethod
-        def Name(cls, number: int) -> str: ...
-        @classmethod
-        def Value(cls, name: str) -> int: ...
-        @classmethod
-        def keys(cls) -> List[str]: ...
-        @classmethod
-        def values(cls) -> List[int]: ...
-        @classmethod
-        def items(cls) -> List[Tuple[str, int]]: ...
-    UNDEFINED = cast(DataType, 0)
-    TENSOR = cast(DataType, 1)
-    SPARSE_TENSOR = cast(DataType, 2)
-    SEQUENCE = cast(DataType, 3)
-    MAP = cast(DataType, 4)
-    OPTIONAL = cast(DataType, 5)
-    
-    name = ... # type: str
-    elem_type = ... # type: int
-    
-    @property
-    def tensor_values(self) -> RepeatedCompositeFieldContainer[TensorProto]: ...
-    
-    @property
-    def sparse_tensor_values(self) -> RepeatedCompositeFieldContainer[SparseTensorProto]: ...
-    
-    @property
-    def sequence_values(self) -> RepeatedCompositeFieldContainer[SequenceProto]: ...
-    
-    @property
-    def map_values(self) -> RepeatedCompositeFieldContainer[MapProto]: ...
-    
-    @property
-    def optional_values(self) -> RepeatedCompositeFieldContainer[OptionalProto]: ...
-    
-    def __init__(self,
-        name : OptionalType[str] = None,
-        elem_type : OptionalType[int] = None,
-        tensor_values : OptionalType[Iterable[TensorProto]] = None,
-        sparse_tensor_values : OptionalType[Iterable[SparseTensorProto]] = None,
-        sequence_values : OptionalType[Iterable[SequenceProto]] = None,
-        map_values : OptionalType[Iterable[MapProto]] = None,
-        optional_values : OptionalType[Iterable[OptionalProto]] = None,
-        ) -> None: ...
-    @classmethod
-    def FromString(cls, s: bytes) -> SequenceProto: ...
-    def MergeFrom(self, other_msg: Message) -> None: ...
-    def CopyFrom(self, other_msg: Message) -> None: ...
-
-class MapProto(Message):
-    name = ... # type: str
-    key_type = ... # type: int
-    keys = ... # type: RepeatedScalarFieldContainer[int]
-    string_keys = ... # type: RepeatedScalarFieldContainer[bytes]
-    
-    @property
-    def values(self) -> SequenceProto: ...
-    
-    def __init__(self,
-        name : OptionalType[str] = None,
-        key_type : OptionalType[int] = None,
-        keys : OptionalType[Iterable[int]] = None,
-        string_keys : OptionalType[Iterable[bytes]] = None,
-        values : OptionalType[SequenceProto] = None,
-        ) -> None: ...
-    @classmethod
-    def FromString(cls, s: bytes) -> MapProto: ...
-    def MergeFrom(self, other_msg: Message) -> None: ...
-    def CopyFrom(self, other_msg: Message) -> None: ...
-
-class OptionalProto(Message):
-    class DataType(int):
-        @classmethod
-        def Name(cls, number: int) -> str: ...
-        @classmethod
-        def Value(cls, name: str) -> int: ...
-        @classmethod
-        def keys(cls) -> List[str]: ...
-        @classmethod
-        def values(cls) -> List[int]: ...
-        @classmethod
-        def items(cls) -> List[Tuple[str, int]]: ...
-    UNDEFINED = cast(DataType, 0)
-    TENSOR = cast(DataType, 1)
-    SPARSE_TENSOR = cast(DataType, 2)
-    SEQUENCE = cast(DataType, 3)
-    MAP = cast(DataType, 4)
-    OPTIONAL = cast(DataType, 5)
-    
-    name = ... # type: str
-    elem_type = ... # type: int
-    
-    @property
-    def tensor_value(self) -> TensorProto: ...
-    
-    @property
-    def sparse_tensor_value(self) -> SparseTensorProto: ...
-    
-    @property
-    def sequence_value(self) -> SequenceProto: ...
-    
-    @property
-    def map_value(self) -> MapProto: ...
-    
-    @property
-    def optional_value(self) -> OptionalProto: ...
-    
-    def __init__(self,
-        name : OptionalType[str] = None,
-        elem_type : OptionalType[int] = None,
-        tensor_value : OptionalType[TensorProto] = None,
-        sparse_tensor_value : OptionalType[SparseTensorProto] = None,
-        sequence_value : OptionalType[SequenceProto] = None,
-        map_value : OptionalType[MapProto] = None,
-        optional_value : OptionalType[OptionalProto] = None,
-        ) -> None: ...
-    @classmethod
-    def FromString(cls, s: bytes) -> OptionalProto: ...
-    def MergeFrom(self, other_msg: Message) -> None: ...
-    def CopyFrom(self, other_msg: Message) -> None: ...
+class OptionalProto(_message.Message):
+    __slots__ = ("name", "elem_type", "tensor_value", "sparse_tensor_value", "sequence_value", "map_value", "optional_value")
+    class DataType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = ()
+        UNDEFINED: _ClassVar[OptionalProto.DataType]
+        TENSOR: _ClassVar[OptionalProto.DataType]
+        SPARSE_TENSOR: _ClassVar[OptionalProto.DataType]
+        SEQUENCE: _ClassVar[OptionalProto.DataType]
+        MAP: _ClassVar[OptionalProto.DataType]
+        OPTIONAL: _ClassVar[OptionalProto.DataType]
+    UNDEFINED: OptionalProto.DataType
+    TENSOR: OptionalProto.DataType
+    SPARSE_TENSOR: OptionalProto.DataType
+    SEQUENCE: OptionalProto.DataType
+    MAP: OptionalProto.DataType
+    OPTIONAL: OptionalProto.DataType
+    NAME_FIELD_NUMBER: _ClassVar[int]
+    ELEM_TYPE_FIELD_NUMBER: _ClassVar[int]
+    TENSOR_VALUE_FIELD_NUMBER: _ClassVar[int]
+    SPARSE_TENSOR_VALUE_FIELD_NUMBER: _ClassVar[int]
+    SEQUENCE_VALUE_FIELD_NUMBER: _ClassVar[int]
+    MAP_VALUE_FIELD_NUMBER: _ClassVar[int]
+    OPTIONAL_VALUE_FIELD_NUMBER: _ClassVar[int]
+    name: str
+    elem_type: int
+    tensor_value: _onnx_ml_pb2.TensorProto
+    sparse_tensor_value: _onnx_ml_pb2.SparseTensorProto
+    sequence_value: SequenceProto
+    map_value: MapProto
+    optional_value: OptionalProto
+    def __init__(self, name: _Optional[str] = ..., elem_type: _Optional[int] = ..., tensor_value: _Optional[_Union[_onnx_ml_pb2.TensorProto, _Mapping]] = ..., sparse_tensor_value: _Optional[_Union[_onnx_ml_pb2.SparseTensorProto, _Mapping]] = ..., sequence_value: _Optional[_Union[SequenceProto, _Mapping]] = ..., map_value: _Optional[_Union[MapProto, _Mapping]] = ..., optional_value: _Optional[_Union[OptionalProto, _Mapping]] = ...) -> None: ...