onnx 1.15.0__cp311-cp311-win_amd64.whl → 1.16.1__cp311-cp311-win_amd64.whl

onnx/helper.py

@@ -42,6 +42,7 @@ from onnx import (
     ValueInfoProto,
     defs,
     mapping,
+    subbyte,
 )
 
 VersionRowType = Union[Tuple[str, int, int, int], Tuple[str, int, int, int, int]]
@@ -74,13 +75,14 @@ VERSION_TABLE: VersionTableType = [
     ("1.14.0", 9, 19, 3, 1),
     ("1.14.1", 9, 19, 3, 1),
     ("1.15.0", 9, 20, 4, 1),
+    ("1.16.0", 10, 21, 5, 1),
 ]
 
 VersionMapType = Dict[Tuple[str, int], int]
 
 
 def create_op_set_id_version_map(table: VersionTableType) -> VersionMapType:
-    """create a map from (opset-domain, opset-version) to ir-version from above table"""
+    """Create a map from (opset-domain, opset-version) to ir-version from above table."""
     result: VersionMapType = {}
 
     def process(release_version: str, ir_version: int, *args: Any) -> None:
@@ -100,13 +102,15 @@ OP_SET_ID_VERSION_MAP = create_op_set_id_version_map(VERSION_TABLE)
 
 
 def find_min_ir_version_for(
-    opsetidlist: List[OperatorSetIdProto], ignore_unknown: bool = False
+    opsetidlist: Sequence[OperatorSetIdProto], ignore_unknown: bool = False
 ) -> int:
     """Given list of opset ids, determine minimum IR version required.
 
-    Arguments:
-        opsetidlist (List[OperatorSetIdProto]): The list of OperatorSetIdProto
-        ignore_unknown (bool): If True, ignore unknown domain and return default min version for that domain.
+    Args:
+        opsetidlist: A sequence of OperatorSetIdProto.
+        ignore_unknown: If True, ignore unknown domain and return default minimum
+            version for that domain.
+
     Returns:
         The minimum IR version required (integer)
     """
@@ -132,11 +136,12 @@ def make_node(
     name: Optional[str] = None,
     doc_string: Optional[str] = None,
     domain: Optional[str] = None,
+    overload: Optional[str] = None,
     **kwargs: Any,
 ) -> NodeProto:
     """Construct a NodeProto.
 
-    Arguments:
+    Args:
         op_type (string): The name of the operator to construct
         inputs (list of string): list of input names
         outputs (list of string): list of output names
@@ -144,12 +149,14 @@ def make_node(
         doc_string (string, default None): optional documentation string for NodeProto
         domain (string, default None): optional domain for NodeProto.
             If it's None, we will just use default domain (which is empty)
+        overload (string, default None): optional field, used to
+            resolve calls to model-local functions
         **kwargs (dict): the attributes of the node.  The acceptable values
             are documented in :func:`make_attribute`.
+
     Returns:
         NodeProto
     """
-
     node = NodeProto()
     node.op_type = op_type
     node.input.extend(inputs)
@@ -160,6 +167,8 @@ def make_node(
         node.doc_string = doc_string
     if domain is not None:
         node.domain = domain
+    if overload is not None:
+        node.overload = overload
     if kwargs:
         node.attribute.extend(
             make_attribute(key, value)
@@ -175,7 +184,7 @@ def make_operatorsetid(
 ) -> OperatorSetIdProto:
     """Construct an OperatorSetIdProto.
 
-    Arguments:
+    Args:
         domain (string): The domain of the operator set id
         version (integer): Version of operator set id
     Returns:
@@ -199,7 +208,7 @@ def make_graph(
 ) -> GraphProto:
     """Construct a GraphProto
 
-    Arguments:
+    Args:
         nodes: list of NodeProto
         name (string): graph name
         inputs: list of ValueInfoProto
@@ -233,7 +242,7 @@ def make_graph(
 def make_opsetid(domain: str, version: int) -> OperatorSetIdProto:
     """Construct an OperatorSetIdProto.
 
-    Arguments:
+    Args:
         domain (string): The domain of the operator set id
         version (integer): Version of operator set id
     Returns:
@@ -255,11 +264,15 @@ def make_function(
     attributes: Optional[Sequence[str]] = None,
     attribute_protos: Optional[Sequence[AttributeProto]] = None,
     doc_string: Optional[str] = None,
+    overload: Optional[str] = None,
+    value_info: Optional[Sequence[ValueInfoProto]] = None,
 ) -> FunctionProto:
     if attributes is None:
         attributes = []
     if attribute_protos is None:
         attribute_protos = []
+    if value_info is None:
+        value_info = []
     f = FunctionProto()
     f.domain = domain
     f.name = fname
@@ -271,13 +284,16 @@ def make_function(
     f.attribute_proto.extend(attribute_protos)
     if doc_string:
         f.doc_string = doc_string
+    if overload is not None:
+        f.overload = overload
+    f.value_info.extend(value_info)
     return f
 
 
 def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
     """Construct a ModelProto
 
-    Arguments:
+    Args:
         graph (GraphProto): *make_graph* returns
         **kwargs: any attribute to add to the returned instance
     Returns:
@@ -315,18 +331,26 @@ def make_model_gen_version(graph: GraphProto, **kwargs: Any) -> ModelProto:
     ir_version_field = "ir_version"
     if ir_version_field not in kwargs:
         opset_imports_field = "opset_imports"
-        imports = kwargs[opset_imports_field] if opset_imports_field in kwargs else []
+        imports = kwargs.get(opset_imports_field, [])
         kwargs[ir_version_field] = find_min_ir_version_for(imports)
     return make_model(graph, **kwargs)
 
 
-def set_model_props(model: ModelProto, dict_value: Dict[str, str]) -> None:
-    del model.metadata_props[:]
+def set_metadata_props(
+    proto: Union[
+        ModelProto, GraphProto, FunctionProto, NodeProto, TensorProto, ValueInfoProto
+    ],
+    dict_value: Dict[str, str],
+) -> None:
+    del proto.metadata_props[:]
     for k, v in dict_value.items():
-        entry = model.metadata_props.add()
+        entry = proto.metadata_props.add()
         entry.key = k
         entry.value = v
-        # model.metadata_properties.append(entry)
+
+
+def set_model_props(model: ModelProto, dict_value: Dict[str, str]) -> None:
+    set_metadata_props(model, dict_value)
 
 
 def split_complex_to_pairs(ca: Sequence[np.complex64]) -> Sequence[int]:
@@ -362,19 +386,22 @@ def float32_to_float8e4m3(  # noqa: PLR0911
     uz: bool = False,
     saturate: bool = True,
 ) -> int:
-    """
-    Convert a float32 value to a float8, e4m3 (as int).
-
-    :param fval: float to convert
-    :param scale: scale, divide *fval* by *scale* before casting it
-    :param fn: no infinite values
-    :param uz: no negative zero
-    :param saturate: if True, any value out of range included inf becomes the maximum value,
-        otherwise, it becomes NaN. The description of operator Cast fully describes the
-        differences.
-    :return: converted float
+    """Convert a float32 value to a float8, e4m3 (as int).
 
     See :ref:`onnx-detail-float8` for technical details.
+
+    Args:
+        fval: float to convert
+        scale: scale, divide *fval* by *scale* before casting it
+        fn: no infinite values
+        uz: no negative zero
+        saturate: if True, any value out of range included inf becomes
+            the maximum value, otherwise, it becomes NaN. The
+            description of operator Cast fully describes the
+            differences.
+
+    Returns:
+        converted float
     """
     if not fn:
         raise NotImplementedError(
@@ -393,47 +420,45 @@ def float32_to_float8e4m3(  # noqa: PLR0911
         e = (b & 0x7F800000) >> 23  # exponent
         m = b & 0x007FFFFF  # mantissa
 
-        if e != 0:
-            if e < 116:  # noqa: PLR2004
-                pass
-            elif e < 120:  # noqa: PLR2004
-                # denormalized number
-                ex = e - 119
-                if ex >= -2:  # noqa: PLR2004
-                    ret |= 1 << (2 + ex)
-                    ret |= m >> (21 - ex)
-                elif m > 0:
-                    ret |= 1
-                mask = 1 << (20 - ex)
-                if m & mask and (
-                    ret & 1
-                    or m & (mask - 1) > 0
-                    or (m & mask and m & (mask << 1) and m & (mask - 1) == 0)
-                ):
+        if e < 116:  # noqa: PLR2004
+            ret = 0
+        elif e < 120:  # noqa: PLR2004
+            # denormalized number
+            ex = e - 119
+            if ex >= -2:  # noqa: PLR2004
+                ret |= 1 << (2 + ex)
+                ret |= m >> (21 - ex)
+            elif m > 0:
+                ret |= 1
+            else:
+                ret = 0
+            mask = 1 << (20 - ex)
+            if m & mask and (
+                ret & 1
+                or m & (mask - 1) > 0
+                or (m & mask and m & (mask << 1) and m & (mask - 1) == 0)
+            ):
+                # rounding
+                ret += 1
+        elif e < 135:  # noqa: PLR2004
+            # normalized number
+            ex = e - 119  # 127 - 8
+            if ex == 0:
+                ret |= 0x4
+                ret |= m >> 21
+            else:
+                ret |= ex << 3
+                ret |= m >> 20
+            if m & 0x80000 and ((m & 0x100000) or (m & 0x7FFFF)):
+                if (ret & 0x7F) < 0x7F:  # noqa: PLR2004
                     # rounding
                     ret += 1
-            elif e < 135:  # noqa: PLR2004
-                # normalized number
-                ex = e - 119  # 127 - 8
-                if ex == 0:
-                    ret |= 0x4
-                    ret |= m >> 21
-                else:
-                    ret |= ex << 3
-                    ret |= m >> 20
-                if m & 0x80000 and ((m & 0x100000) or (m & 0x7FFFF)):
-                    if (ret & 0x7F) < 0x7F:  # noqa: PLR2004
-                        # rounding
-                        ret += 1
-                    elif not saturate:
-                        return 0x80
-            elif saturate:
-                ret |= 0x7F  # 01111110
-            else:
-                ret = 0x80
-        elif m == 0:
-            # -0
-            ret = 0
+                elif not saturate:
+                    return 0x80
+        elif saturate:
+            ret |= 0x7F  # 01111110
+        else:
+            ret = 0x80
         return int(ret)
     else:
         if (b & 0x7FC00000) == 0x7FC00000:  # noqa: PLR2004
@@ -495,17 +520,20 @@ def float32_to_float8e5m2(  # noqa: PLR0911
     uz: bool = False,
     saturate: bool = True,
 ) -> int:
-    """
-    Convert a float32 value to a float8, e5m2 (as int).
-
-    :param fval: float to convert
-    :param scale: scale, divide *fval* by *scale* before casting it
-    :param fn: no infinite values
-    :param uz: no negative zero
-    :param saturate: if True, any value out of range included inf becomes the maximum value,
-        otherwise, it becomes NaN. The description of operator Cast fully describes the
-        differences.
-    :return: converted float
+    """Convert a float32 value to a float8, e5m2 (as int).
+
+    Args:
+        fval: float to convert
+        scale: scale, divide *fval* by *scale* before casting it
+        fn: no infinite values
+        uz: no negative zero
+        saturate: if True, any value out of range included inf becomes
+            the maximum value, otherwise, it becomes NaN. The
+            description of operator Cast fully describes the
+            differences.
+
+    Returns:
+        converted float
     """
     x = fval / scale
     b = int.from_bytes(struct.pack("<f", np.float32(x)), "little")
@@ -522,45 +550,43 @@ def float32_to_float8e5m2(  # noqa: PLR0911
         e = (b & 0x7F800000) >> 23  # exponent
         m = b & 0x007FFFFF  # mantissa
 
-        if e != 0:
-            if e < 109:  # noqa: PLR2004
-                pass
-            elif e < 112:  # noqa: PLR2004
-                # denormalized number
-                ex = e - 111
-                if ex >= -1:
-                    ret |= 1 << (1 + ex)
-                    ret |= m >> (22 - ex)
-                elif m > 0:
-                    ret |= 1
-                mask = 1 << (21 - ex)
-                if m & mask and (
-                    ret & 1
-                    or m & (mask - 1) > 0
-                    or (m & mask and m & (mask << 1) and m & (mask - 1) == 0)
-                ):
+        if e < 109:  # noqa: PLR2004
+            ret = 0
+        elif e < 112:  # noqa: PLR2004
+            # denormalized number
+            ex = e - 111
+            if ex >= -1:
+                ret |= 1 << (1 + ex)
+                ret |= m >> (22 - ex)
+            elif m > 0:
+                ret |= 1
+            else:
+                ret = 0
+            mask = 1 << (21 - ex)
+            if m & mask and (
+                ret & 1
+                or m & (mask - 1) > 0
+                or (m & mask and m & (mask << 1) and m & (mask - 1) == 0)
+            ):
+                # rounding
+                ret += 1
+        elif e < 143:  # noqa: PLR2004
+            # normalized number
+            ex = e - 111
+            ret |= ex << 2
+            ret |= m >> 21
+            if m & 0x100000 and ((m & 0xFFFFF) or (m & 0x200000)):
+                if (ret & 0x7F) < 0x7F:  # noqa: PLR2004
                     # rounding
                     ret += 1
-            elif e < 143:  # noqa: PLR2004
-                # normalized number
-                ex = e - 111
-                ret |= ex << 2
-                ret |= m >> 21
-                if m & 0x100000 and ((m & 0xFFFFF) or (m & 0x200000)):
-                    if (ret & 0x7F) < 0x7F:  # noqa: PLR2004
-                        # rounding
-                        ret += 1
-                    elif not saturate:
-                        ret = 0x80
-            elif e == 255 and m == 0:  # inf  # noqa: PLR2004
-                ret = 0x80
-            elif saturate:
-                ret |= 0x7F  # last possible number
-            else:
-                ret = 0x80
-        elif m == 0:
-            # -0
-            ret = 0
+                elif not saturate:
+                    ret = 0x80
+        elif e == 255 and m == 0:  # inf  # noqa: PLR2004
+            ret = 0x80
+        elif saturate:
+            ret |= 0x7F  # last possible number
+        else:
+            ret = 0x80
         return int(ret)
     elif not fn and not uz:
         if (b & 0x7FC00000) == 0x7FC00000:  # noqa: PLR2004
@@ -613,17 +639,44 @@ def float32_to_float8e5m2(  # noqa: PLR0911
         raise NotImplementedError("fn and uz must be both False or True.")
 
 
+def pack_float32_to_4bit(
+    array: Union[np.ndarray, Sequence], signed: bool
+) -> np.ndarray:
+    """Convert an array of float32 value to a 4bit data-type and pack every two concecutive elements in a byte.
+    See :ref:`onnx-detail-int4` for technical details.
+
+    Args:
+        array: array of float to convert and pack
+        signed: Whether the 4 bit variant is signed or unsigned
+
+    Returns:
+        Packed array with size `ceil(farray.size/2)` (single dimension).
+    """
+    if not isinstance(array, np.ndarray):
+        array = np.asarray(array, dtype=np.float32)
+
+    array_flat = array.ravel()
+    is_odd_volume = np.prod(array.shape) % 2 == 1
+    if is_odd_volume:
+        array_flat = np.append(array_flat, np.array([0]))
+
+    single_func = lambda x, y: subbyte.float32x2_to_4bitx2(x, y, signed)  # noqa: E731
+    func = np.frompyfunc(single_func, 2, 1)
+
+    arr = func(array_flat[0::2], array_flat[1::2])
+    return arr.astype(np.uint8)  # type: ignore[no-any-return]
+
+
 def make_tensor(
     name: str, data_type: int, dims: Sequence[int], vals: Any, raw: bool = False
 ) -> TensorProto:
-    """
-    Make a TensorProto with specified arguments.  If raw is False, this
+    """Make a TensorProto with specified arguments.  If raw is False, this
     function will choose the corresponding proto field to store the
     values based on data_type. If raw is True, use "raw_data" proto
     field to store the values, and values should be of type bytes in
     this case.
 
-    Arguments:
+    Args:
         name (string): tensor name
         data_type (int): a value such as onnx.TensorProto.FLOAT
         dims (List[int]): shape
@@ -646,8 +699,7 @@ def make_tensor(
     # Check number of vals specified equals tensor size
     expected_size = 1
     if raw:
-        # NumPy doesn't have BFLOAT16. TENSOR_TYPE_TO_NP_TYPE maps it to float32,
-        # which has the wrong itemsize.
+        # NumPy doesn't have BFLOAT16. TENSOR_TYPE_MAP maps it to float32, which has the wrong itemsize.
         if data_type == TensorProto.BFLOAT16:
             expected_size = 2
         elif data_type in (
@@ -657,6 +709,9 @@ def make_tensor(
             TensorProto.FLOAT8E5M2FNUZ,
         ):
             expected_size = 1
+        # NumPy doesn't have INT4. It is packed in couples to UINT8 buffers.
+        elif data_type in (TensorProto.UINT4, TensorProto.INT4):
+            expected_size = 0.5  # type: ignore[assignment]
         else:
             expected_size = np_dtype.itemsize
 
@@ -666,9 +721,14 @@ def make_tensor(
         expected_size *= d
 
     if len(vals) != expected_size:
-        raise ValueError(
-            f"Number of values does not match tensor's size. Expected {expected_size}, but it is {len(vals)}. "
-        )
+        # padding of half a byte is acceptable for 4bit types
+        if not (
+            data_type in (TensorProto.UINT4, TensorProto.INT4)
+            and len(vals) == expected_size + 0.5
+        ):
+            raise ValueError(
+                f"Number of values does not match tensor's size. Expected {expected_size}, but it is {len(vals)}. "
+            )
 
     if raw:
         tensor.raw_data = vals
@@ -705,6 +765,17 @@ def make_tensor(
                     np.array(vals).astype(np_dtype).flatten().tolist(),
                 )
             )
+        elif data_type in (
+            TensorProto.UINT4,
+            TensorProto.INT4,
+        ):
+            signed = data_type == TensorProto.INT4
+            vals = (
+                pack_float32_to_4bit(vals, signed=signed)
+                .astype(np_dtype)
+                .flatten()
+                .tolist()
+            )
         elif data_type == TensorProto.BOOL:
             vals = np.array(vals).astype(int)
         elif data_type == TensorProto.STRING:
@@ -720,7 +791,7 @@ def make_sparse_tensor(
 ) -> SparseTensorProto:
     """Construct a SparseTensorProto
 
-    Arguments:
+    Args:
         values (TensorProto): the values
         indices (TensorProto): the indices
         dims: the shape
@@ -740,9 +811,7 @@ def make_sequence(
     elem_type: SequenceProto.DataType,
     values: Sequence[Any],
 ) -> SequenceProto:
-    """
-    Make a Sequence with specified value arguments.
-    """
+    """Make a Sequence with specified value arguments."""
     sequence = SequenceProto()
     sequence.name = name
     sequence.elem_type = elem_type
@@ -769,8 +838,7 @@ def make_sequence(
 def make_map(
     name: str, key_type: int, keys: List[Any], values: SequenceProto
 ) -> MapProto:
-    """
-    Make a Map with specified key-value pair arguments.
+    """Make a Map with specified key-value pair arguments.
 
     Criteria for conversion:
     - Keys and Values must have the same number of elements
@@ -803,9 +871,7 @@ def make_optional(
     elem_type: OptionalProto.DataType,
     value: Optional[Any],
 ) -> OptionalProto:
-    """
-    Make an Optional with specified value arguments.
-    """
+    """Make an Optional with specified value arguments."""
     optional = OptionalProto()
     optional.name = name
     optional.elem_type = elem_type
@@ -997,7 +1063,6 @@ def make_tensor_type_proto(
     shape_denotation: Optional[List[str]] = None,
 ) -> TypeProto:
     """Makes a Tensor TypeProto based on the data type and shape."""
-
     type_proto = TypeProto()
     tensor_type_proto = type_proto.tensor_type
     tensor_type_proto.elem_type = elem_type
@@ -1061,7 +1126,6 @@ def make_sparse_tensor_type_proto(
     shape_denotation: Optional[List[str]] = None,
 ) -> TypeProto:
     """Makes a SparseTensor TypeProto based on the data type and shape."""
-
     type_proto = TypeProto()
     sparse_tensor_type_proto = type_proto.sparse_tensor_type
     sparse_tensor_type_proto.elem_type = elem_type
@@ -1290,7 +1354,7 @@ def printable_type(t: TypeProto) -> str:
                 s += str(", " + "x".join(map(printable_dim, t.tensor_type.shape.dim)))
             else:
                 s += ", scalar"
-        return s
+        return s  # type: ignore[no-any-return]
     if t.WhichOneof("value") is None:
         return ""
     return f"Unknown type {t.WhichOneof('value')}"
@@ -1351,10 +1415,9 @@ def printable_node(
 
 
 def printable_graph(graph: GraphProto, prefix: str = "") -> str:
-    """
-    Display a GraphProto as a string.
+    """Display a GraphProto as a string.
 
-    Arguments:
+    Args:
         graph (GraphProto): the graph to display
         prefix (string): prefix of every line
 
@@ -1431,9 +1494,7 @@ def printable_graph(graph: GraphProto, prefix: str = "") -> str:
 
 
 def strip_doc_string(proto: google.protobuf.message.Message) -> None:
-    """
-    Empties `doc_string` field on any nested protobuf messages
-    """
+    """Empties `doc_string` field on any nested protobuf messages"""
     if not isinstance(proto, google.protobuf.message.Message):
         raise TypeError(
             f"proto must be an instance of {google.protobuf.message.Message}."
@@ -1475,41 +1536,49 @@ def make_training_info(
 
 # Following functions are used for mapping
 def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
-    """
-    Convert a TensorProto's data_type to corresponding numpy dtype. It can be used while making tensor.
+    """Convert a TensorProto's data_type to corresponding numpy dtype. It can be used while making tensor.
+
+    Args:
+        tensor_dtype: TensorProto's data_type
 
-    :param tensor_dtype: TensorProto's data_type
-    :return: numpy's data_type
+    Returns:
+        numpy's data_type
     """
     return mapping.TENSOR_TYPE_MAP[tensor_dtype].np_dtype
 
 
 def tensor_dtype_to_storage_tensor_dtype(tensor_dtype: int) -> int:
-    """
-    Convert a TensorProto's data_type to corresponding data_type for storage.
+    """Convert a TensorProto's data_type to corresponding data_type for storage.
 
-    :param tensor_dtype: TensorProto's data_type
-    :return: data_type for storage
+    Args:
+        tensor_dtype: TensorProto's data_type
+
+    Returns:
+        data_type for storage
     """
     return mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
 
 
 def tensor_dtype_to_string(tensor_dtype: int) -> str:
-    """
-    Get the name of given TensorProto's data_type.
+    """Get the name of given TensorProto's data_type.
+
+    Args:
+        tensor_dtype: TensorProto's data_type
 
-    :param tensor_dtype: TensorProto's data_type
-    :return: the name of data_type
+    Returns:
+        the name of data_type
     """
     return mapping.TENSOR_TYPE_MAP[tensor_dtype].name
 
 
 def tensor_dtype_to_field(tensor_dtype: int) -> str:
-    """
-    Convert a TensorProto's data_type to corresponding field name for storage. It can be used while making tensors.
+    """Convert a TensorProto's data_type to corresponding field name for storage. It can be used while making tensors.
 
-    :param tensor_dtype: TensorProto's data_type
-    :return: field name
+    Args:
+        tensor_dtype: TensorProto's data_type
+
+    Returns:
+        field name
     """
     return mapping._STORAGE_TENSOR_TYPE_TO_FIELD[
         mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
@@ -1517,11 +1586,13 @@ def tensor_dtype_to_field(tensor_dtype: int) -> str:
 
 
 def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
-    """
-    Convert a numpy's dtype to corresponding tensor type. It can be used while converting numpy arrays to tensors.
+    """Convert a numpy's dtype to corresponding tensor type. It can be used while converting numpy arrays to tensors.
+
+    Args:
+        np_dtype: numpy's data_type
 
-    :param np_dtype: numpy's data_type
-    :return: TensorsProto's data_type
+    Returns:
+        TensorsProto's data_type
     """
     return cast(
         int,
@@ -1530,10 +1601,10 @@ def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> int:
 
 
 def get_all_tensor_dtypes() -> KeysView[int]:
-    """
-    Get all tensor types from TensorProto.
+    """Get all tensor types from TensorProto.
 
-    :return: all tensor types from TensorProto
+    Returns:
+        all tensor types from TensorProto
     """
     return mapping.TENSOR_TYPE_MAP.keys()
 
@@ -1542,12 +1613,14 @@ _ATTRIBUTE_TYPE_TO_STR = {k: v for v, k in AttributeProto.AttributeType.items()}
 
 
 def _attr_type_to_str(attr_type: int) -> str:
-    """
-    Convert AttributeProto type to string.
+    """Convert AttributeProto type to string.
 
-    :param attr_type: AttributeProto type.
-    :return: String representing the supplied attr_type.
+    Args:
+        attr_type: AttributeProto type.
+
+    Returns:
+        String representing the supplied attr_type.
     """
     if attr_type in AttributeProto.AttributeType.values():
-        return _ATTRIBUTE_TYPE_TO_STR[attr_type]
-    return AttributeProto.AttributeType.keys()[0]
+        return _ATTRIBUTE_TYPE_TO_STR[attr_type]  # type: ignore[no-any-return]
+    return AttributeProto.AttributeType.keys()[0]  # type: ignore[no-any-return]