Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/onnx/helper.py: 28%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (c) ONNX Project Contributors
2#
3# SPDX-License-Identifier: Apache-2.0
4from __future__ import annotations
6import collections.abc
7import functools
8import math
9import numbers
10import typing
11from typing import TYPE_CHECKING, Any, TypeVar
13import google.protobuf.message
14import numpy as np
15import typing_extensions
17import onnx
18from onnx import _mapping, defs
19from onnx.onnx_data_pb import MapProto, OptionalProto, SequenceProto
20from onnx.onnx_pb import (
21 AttributeProto,
22 FunctionProto,
23 GraphProto,
24 ModelProto,
25 NodeProto,
26 OperatorSetIdProto,
27 TensorProto,
28 TensorShapeProto,
29 TrainingInfoProto,
30 TypeProto,
31 ValueInfoProto,
32)
34if TYPE_CHECKING:
35 from collections.abc import Callable, KeysView, Sequence
37 from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
39VersionRowType = tuple[str, int, int, int] | tuple[str, int, int, int, int]
40VersionTableType = list[VersionRowType]
41AssignmentBindingType = list[tuple[str, str]]
43# This is a copy of the documented version in https://github.com/onnx/onnx/blob/main/docs/Versioning.md#released-versions
44# Both must be updated whenever a new version of ONNX is released.
45VERSION_TABLE: VersionTableType = [
46 # Release-version, IR version, ai.onnx version, ai.onnx.ml version, (optional) ai.onnx.training version
47 ("1.0", 3, 1, 1),
48 ("1.1", 3, 5, 1),
49 ("1.1.2", 3, 6, 1),
50 ("1.2", 3, 7, 1),
51 ("1.3", 3, 8, 1),
52 ("1.4.1", 4, 9, 1),
53 ("1.5.0", 5, 10, 1),
54 ("1.6.0", 6, 11, 2),
55 ("1.7.0", 7, 12, 2, 1),
56 ("1.8.0", 7, 13, 2, 1),
57 ("1.8.1", 7, 13, 2, 1),
58 ("1.9.0", 7, 14, 2, 1),
59 ("1.10.0", 8, 15, 2, 1),
60 ("1.10.1", 8, 15, 2, 1),
61 ("1.10.2", 8, 15, 2, 1),
62 ("1.11.0", 8, 16, 3, 1),
63 ("1.12.0", 8, 17, 3, 1),
64 ("1.13.0", 8, 18, 3, 1),
65 ("1.13.1", 8, 18, 3, 1),
66 ("1.14.0", 9, 19, 3, 1),
67 ("1.14.1", 9, 19, 3, 1),
68 ("1.15.0", 9, 20, 4, 1),
69 ("1.16.0", 10, 21, 5, 1),
70 ("1.16.1", 10, 21, 5, 1),
71 ("1.16.2", 10, 21, 5, 1),
72 ("1.17.0", 10, 22, 5, 1),
73 ("1.18.0", 11, 23, 5, 1),
74 ("1.19.0", 12, 24, 5, 1),
75 ("1.19.1", 12, 24, 5, 1),
76 ("1.20.0", 13, 25, 5, 1),
77 ("1.20.1", 13, 25, 5, 1),
78 ("1.21.0", 13, 26, 5, 1),
79 ("1.22.0", 13, 27, 5, 1),
80 ("1.23.0", 14, 28, 5, 1),
81]
83VersionMapType = dict[tuple[str, int], int]
86def _create_op_set_id_version_map(table: VersionTableType) -> VersionMapType:
87 """Create a map from (opset-domain, opset-version) to ir-version from above table."""
88 result: VersionMapType = {}
89 for row in table:
90 ir_version = row[1]
91 for pair in zip(
92 ["ai.onnx", "ai.onnx.ml", "ai.onnx.training"],
93 row[2:],
94 strict=False,
95 ):
96 if pair not in result:
97 result[pair] = ir_version
98 if pair[0] == "ai.onnx":
99 result["ai.onnx.preview", pair[1]] = ir_version
100 if pair[0] == "ai.onnx.training":
101 result["ai.onnx.preview.training", pair[1]] = ir_version
102 return result
105OP_SET_ID_VERSION_MAP = _create_op_set_id_version_map(VERSION_TABLE)
108def find_min_ir_version_for(
109 opsetidlist: Sequence[OperatorSetIdProto], ignore_unknown: bool = False
110) -> int:
111 """Given list of opset ids, determine minimum IR version required.
113 Args:
114 opsetidlist: A sequence of OperatorSetIdProto.
115 ignore_unknown: If True, ignore unknown domain and return default minimum
116 version for that domain.
118 Returns:
119 The minimum IR version required (integer)
120 """
121 default_min_version = 3
123 def find_min(domain: str | None, version: int) -> int:
124 key = (domain or "ai.onnx", version)
125 if key in OP_SET_ID_VERSION_MAP:
126 return OP_SET_ID_VERSION_MAP[key]
127 if ignore_unknown:
128 return default_min_version
129 raise ValueError("Unsupported opset-version.")
131 if opsetidlist:
132 return max(find_min(x.domain, x.version) for x in opsetidlist)
133 return default_min_version # if no opsets specified
136def make_node(
137 op_type: str,
138 inputs: Sequence[str],
139 outputs: Sequence[str],
140 name: str | None = None,
141 doc_string: str | None = None,
142 domain: str | None = None,
143 overload: str | None = None,
144 **kwargs: Any,
145) -> NodeProto:
146 """Construct a NodeProto.
148 Args:
149 op_type (string): The name of the operator to construct
150 inputs (list of string): list of input names
151 outputs (list of string): list of output names
152 name (string, default None): optional unique identifier for NodeProto
153 doc_string (string, default None): optional documentation string for NodeProto
154 domain (string, default None): optional domain for NodeProto.
155 If it's None, we will just use default domain (which is empty)
156 overload (string, default None): optional field, used to
157 resolve calls to model-local functions
158 **kwargs (dict): the attributes of the node. The acceptable values
159 are documented in :func:`make_attribute`.
161 Returns:
162 NodeProto
163 """
164 node = NodeProto()
165 node.op_type = op_type
166 node.input.extend(inputs)
167 node.output.extend(outputs)
168 if name:
169 node.name = name
170 if doc_string:
171 node.doc_string = doc_string
172 if domain is not None:
173 node.domain = domain
174 if overload is not None:
175 node.overload = overload
176 if kwargs:
177 node.attribute.extend(
178 make_attribute(key, value)
179 for key, value in sorted(kwargs.items())
180 if value is not None
181 )
182 return node
185def make_operatorsetid(
186 domain: str,
187 version: int,
188) -> OperatorSetIdProto:
189 """Construct an OperatorSetIdProto.
191 Args:
192 domain (string): The domain of the operator set id
193 version (integer): Version of operator set id
194 Returns:
195 OperatorSetIdProto
196 """
197 operatorsetid = OperatorSetIdProto()
198 operatorsetid.domain = domain
199 operatorsetid.version = version
200 return operatorsetid
203def make_graph(
204 nodes: Sequence[NodeProto],
205 name: str,
206 inputs: Sequence[ValueInfoProto],
207 outputs: Sequence[ValueInfoProto],
208 initializer: Sequence[TensorProto] | None = None,
209 doc_string: str | None = None,
210 value_info: Sequence[ValueInfoProto] | None = None,
211 sparse_initializer: Sequence[onnx.SparseTensorProto] | None = None,
212) -> GraphProto:
213 """Construct a GraphProto
215 Args:
216 nodes: list of NodeProto
217 name (string): graph name
218 inputs: list of ValueInfoProto
219 outputs: list of ValueInfoProto
220 initializer: list of TensorProto
221 doc_string (string): graph documentation
222 value_info: list of ValueInfoProto
223 sparse_initializer: list of onnx.SparseTensorProto
224 Returns:
225 GraphProto
226 """
227 if initializer is None:
228 initializer = []
229 if sparse_initializer is None:
230 sparse_initializer = []
231 if value_info is None:
232 value_info = []
233 graph = GraphProto()
234 graph.node.extend(nodes)
235 graph.name = name
236 graph.input.extend(inputs)
237 graph.output.extend(outputs)
238 graph.initializer.extend(initializer)
239 graph.sparse_initializer.extend(sparse_initializer)
240 graph.value_info.extend(value_info)
241 if doc_string:
242 graph.doc_string = doc_string
243 return graph
246def make_opsetid(domain: str, version: int) -> OperatorSetIdProto:
247 """Construct an OperatorSetIdProto.
249 Args:
250 domain (string): The domain of the operator set id
251 version (integer): Version of operator set id
252 Returns:
253 OperatorSetIdProto
254 """
255 opsetid = OperatorSetIdProto()
256 opsetid.domain = domain
257 opsetid.version = version
258 return opsetid
261def make_function(
262 domain: str,
263 fname: str,
264 inputs: Sequence[str],
265 outputs: Sequence[str],
266 nodes: Sequence[NodeProto],
267 opset_imports: Sequence[OperatorSetIdProto],
268 attributes: Sequence[str] | None = None,
269 attribute_protos: Sequence[AttributeProto] | None = None,
270 doc_string: str | None = None,
271 overload: str | None = None,
272 value_info: Sequence[ValueInfoProto] | None = None,
273) -> FunctionProto:
274 if attributes is None:
275 attributes = []
276 if attribute_protos is None:
277 attribute_protos = []
278 if value_info is None:
279 value_info = []
280 f = FunctionProto()
281 f.domain = domain
282 f.name = fname
283 f.input.extend(inputs)
284 f.output.extend(outputs)
285 f.node.extend(nodes)
286 f.opset_import.extend(opset_imports)
287 f.attribute.extend(attributes)
288 f.attribute_proto.extend(attribute_protos)
289 if doc_string:
290 f.doc_string = doc_string
291 if overload is not None:
292 f.overload = overload
293 f.value_info.extend(value_info)
294 return f
297def make_model(graph: GraphProto, **kwargs: Any) -> ModelProto:
298 """Construct a ModelProto
300 Args:
301 graph (GraphProto): *make_graph* returns
302 **kwargs: any attribute to add to the returned instance
303 Returns:
304 ModelProto
305 """
306 model = ModelProto()
307 # Touch model.ir_version so it is stored as the version from which it is
308 # generated.
309 model.ir_version = onnx.IR_VERSION
310 model.graph.CopyFrom(graph)
312 opset_imports: Sequence[OperatorSetIdProto] | None = kwargs.pop(
313 "opset_imports", None
314 )
315 if opset_imports is not None:
316 model.opset_import.extend(opset_imports)
317 else:
318 # Default import
319 imp = model.opset_import.add()
320 imp.version = defs.onnx_opset_version()
322 functions: Sequence[FunctionProto] | None = kwargs.pop("functions", None)
323 if functions is not None:
324 model.functions.extend(functions)
326 for k, v in kwargs.items():
327 # TODO: Does this work with repeated fields?
328 setattr(model, k, v)
329 return model
332# An extension of make_model that infers an IR_VERSION for the model,
333# if not specified, using a best-effort-basis.
334def make_model_gen_version(graph: GraphProto, **kwargs: Any) -> ModelProto:
335 ir_version_field = "ir_version"
336 if ir_version_field not in kwargs:
337 opset_imports_field = "opset_imports"
338 imports = kwargs.get(opset_imports_field, [])
339 kwargs[ir_version_field] = find_min_ir_version_for(imports)
340 return make_model(graph, **kwargs)
343def set_metadata_props(
344 proto: (
345 ModelProto
346 | GraphProto
347 | FunctionProto
348 | NodeProto
349 | TensorProto
350 | ValueInfoProto
351 ),
352 dict_value: dict[str, str],
353) -> None:
354 del proto.metadata_props[:]
355 for k, v in dict_value.items():
356 entry = proto.metadata_props.add()
357 entry.key = k
358 entry.value = v
361def set_model_props(model: ModelProto, dict_value: dict[str, str]) -> None:
362 set_metadata_props(model, dict_value)
365def make_tensor(
366 name: str,
367 data_type: int,
368 dims: Sequence[int],
369 vals: Sequence[int | float] | bytes | np.ndarray,
370 raw: bool = False,
371) -> TensorProto:
372 """Make a TensorProto with specified arguments. If raw is False, this
373 function will choose the corresponding proto field to store the
374 values based on data_type. If raw is True, use "raw_data" proto
375 field to store the values, and values should be of type bytes in
376 this case.
378 Args:
379 name: tensor name
380 data_type: a value such as onnx.TensorProto.FLOAT
381 dims: shape
382 vals: values
383 raw: if True, vals contains the serialized content of the tensor,
384 otherwise, vals should be a list of values of the type defined by ``data_type``.
386 Returns:
387 TensorProto
388 """
389 tensor = TensorProto()
390 tensor.data_type = data_type
391 tensor.name = name
392 tensor.dims.extend(dims)
394 if data_type == TensorProto.STRING and raw:
395 raise TypeError("Can not use raw_data to store string type.")
397 np_dtype = tensor_dtype_to_np_dtype(data_type)
399 if raw:
400 # NumPy doesn't have INT2/INT4/FP4/FP6. It is packed in couples to UINT8 buffers.
401 if data_type in {TensorProto.FLOAT6E2M3, TensorProto.FLOAT6E3M2}:
402 bits_per_element = 6
403 elif data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
404 bits_per_element = 4
405 elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
406 bits_per_element = 2
407 else:
408 bits_per_element = np_dtype.itemsize * 8
409 # Integer ceil(n * bits_per_element / 8): avoids the float-multiplication
410 # precision loss for large n that a `math.ceil(fraction * n)` formula
411 # would be susceptible to.
412 expected_size_bytes = -(-math.prod(dims) * bits_per_element // 8)
413 if isinstance(vals, np.ndarray):
414 if data_type in {
415 TensorProto.INT4,
416 TensorProto.UINT4,
417 TensorProto.FLOAT4E2M1,
418 }:
419 vals = onnx.numpy_helper._pack_4bitx2(vals)
420 elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
421 vals = onnx.numpy_helper._pack_2bitx4(vals)
422 elif data_type in {TensorProto.FLOAT6E2M3, TensorProto.FLOAT6E3M2}:
423 vals = onnx.numpy_helper._pack_6bit(vals.view(np.uint8))
425 raw_data = onnx.numpy_helper.tobytes_little_endian(vals)
426 elif isinstance(vals, bytes):
427 raw_data = vals
428 else:
429 raise TypeError(
430 f"Raw data must be bytes or numpy.ndarray, but got {type(vals)}."
431 )
432 if len(raw_data) != expected_size_bytes:
433 raise ValueError(
434 f"Raw data size does not match tensor's size. Expected {expected_size_bytes} bytes, but got {len(raw_data)} bytes."
435 )
436 tensor.raw_data = raw_data
437 return tensor
439 assert not raw, "Bug: raw should be False at this point."
441 if data_type == TensorProto.STRING:
442 vals = np.array(vals).flatten()
443 if len(vals) != 0:
444 vals = np.vectorize(_to_bytes)(vals) # Convert to bytes
445 elif data_type in {
446 TensorProto.FLOAT8E4M3FN,
447 TensorProto.FLOAT8E4M3FNUZ,
448 TensorProto.FLOAT8E5M2,
449 TensorProto.FLOAT8E5M2FNUZ,
450 }:
451 # Float8 values are by default casted using saturating cast.
452 vals = onnx.numpy_helper.saturate_cast(np.asarray(vals), np_dtype).flatten()
453 elif data_type == TensorProto.FLOAT8E8M0:
454 vals = onnx.numpy_helper.to_float8e8m0(
455 np.asarray(vals), saturate=True, round_mode="up"
456 ).flatten()
457 elif data_type in {TensorProto.FLOAT6E2M3, TensorProto.FLOAT6E3M2}:
458 # For FP6 in non-raw mode, store per-element bytes in int32_data (like float8)
459 vals = np.asarray(vals, dtype=np_dtype).flatten().view(np.uint8)
460 else:
461 vals = np.asarray(vals, dtype=np_dtype).flatten()
463 expected_elements = math.prod(dims)
464 if len(vals) != expected_elements:
465 raise ValueError(
466 f"Number of values ({len(vals)}) does not match tensor "
467 f"dimensions requiring {expected_elements} elements."
468 )
469 if data_type == TensorProto.COMPLEX128:
470 vals = vals.view(np.float64) # type: ignore[union-attr]
471 elif data_type == TensorProto.COMPLEX64:
472 vals = vals.view(np.float32) # type: ignore[union-attr]
473 elif data_type in {TensorProto.BFLOAT16, TensorProto.FLOAT16}:
474 vals = vals.view(np.uint16) # type: ignore[union-attr]
475 elif data_type in {
476 TensorProto.FLOAT8E4M3FN,
477 TensorProto.FLOAT8E4M3FNUZ,
478 TensorProto.FLOAT8E5M2,
479 TensorProto.FLOAT8E5M2FNUZ,
480 TensorProto.FLOAT8E8M0,
481 }:
482 vals = vals.view(np.uint8) # type: ignore[union-attr]
483 elif data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
484 # Convert to packed 4-bit representation
485 vals = onnx.numpy_helper._pack_4bitx2(vals) # type: ignore[arg-type]
486 elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
487 # Convert to packed 2-bit representation
488 vals = onnx.numpy_helper._pack_2bitx4(vals) # type: ignore[arg-type]
489 elif data_type == TensorProto.BOOL:
490 vals = vals.astype(np.uint8) # type: ignore[union-attr]
492 field = tensor_dtype_to_field(data_type)
493 getattr(tensor, field).extend(vals)
494 return tensor
497def make_sparse_tensor(
498 values: TensorProto, indices: TensorProto, dims: Sequence[int]
499) -> onnx.SparseTensorProto:
500 """Construct a SparseTensorProto
502 Args:
503 values (TensorProto): the values
504 indices (TensorProto): the indices
505 dims: the shape
507 Returns:
508 SparseTensorProto
509 """
510 sparse = onnx.SparseTensorProto()
511 sparse.values.CopyFrom(values)
512 sparse.indices.CopyFrom(indices)
513 sparse.dims.extend(dims)
514 return sparse
517def make_sequence(
518 name: str,
519 elem_type: SequenceProto.DataType,
520 values: Sequence[Any],
521) -> SequenceProto:
522 """Make a Sequence with specified value arguments."""
523 sequence = SequenceProto()
524 sequence.name = name
525 sequence.elem_type = elem_type # type: ignore[assignment]
527 if elem_type == SequenceProto.UNDEFINED:
528 return sequence
530 attribute: RepeatedCompositeFieldContainer | None = None
531 if elem_type == SequenceProto.TENSOR:
532 attribute = sequence.tensor_values
533 elif elem_type == SequenceProto.SPARSE_TENSOR:
534 attribute = sequence.sparse_tensor_values
535 elif elem_type == SequenceProto.SEQUENCE:
536 attribute = sequence.sequence_values
537 elif elem_type == SequenceProto.MAP:
538 attribute = sequence.map_values
539 elif elem_type == OptionalProto.OPTIONAL:
540 attribute = sequence.optional_values
541 else:
542 raise TypeError("The element type in the input sequence is not supported.")
544 attribute.extend(values)
545 return sequence
548def make_map(
549 name: str, key_type: int, keys: list[Any], values: SequenceProto
550) -> MapProto:
551 """Make a Map with specified key-value pair arguments.
553 Criteria for conversion:
554 - Keys and Values must have the same number of elements
555 - Every key in keys must be of the same type
556 - Every value in values must be of the same type
557 """
558 map_proto = MapProto()
559 valid_key_int_types = [
560 TensorProto.INT8,
561 TensorProto.INT16,
562 TensorProto.INT32,
563 TensorProto.INT64,
564 TensorProto.UINT8,
565 TensorProto.UINT16,
566 TensorProto.UINT32,
567 TensorProto.UINT64,
568 ]
569 map_proto.name = name
570 map_proto.key_type = key_type
571 if key_type == TensorProto.STRING:
572 map_proto.string_keys.extend(keys)
573 elif key_type in valid_key_int_types:
574 map_proto.keys.extend(keys)
575 map_proto.values.CopyFrom(values)
576 return map_proto
579def make_optional(
580 name: str,
581 elem_type: OptionalProto.DataType,
582 value: google.protobuf.message.Message | None,
583) -> OptionalProto:
584 """Make an Optional with specified value arguments."""
585 optional = OptionalProto()
586 optional.name = name
587 optional.elem_type = elem_type # type: ignore[assignment]
589 if elem_type == OptionalProto.UNDEFINED:
590 return optional
591 attribute: google.protobuf.message.Message | None = None
592 if elem_type == OptionalProto.TENSOR:
593 attribute = optional.tensor_value
594 elif elem_type == OptionalProto.SPARSE_TENSOR:
595 attribute = optional.sparse_tensor_value
596 elif elem_type == OptionalProto.SEQUENCE:
597 attribute = optional.sequence_value
598 elif elem_type == OptionalProto.MAP:
599 attribute = optional.map_value
600 elif elem_type == OptionalProto.OPTIONAL:
601 attribute = optional.optional_value
602 else:
603 raise TypeError("The element type in the input optional is not supported.")
605 assert value is not None
606 attribute.CopyFrom(value) # type: ignore[arg-type]
607 return optional
610def _to_bytes(value: str | bytes) -> bytes:
611 """Coerce a string (or bytes) value into UTF-8 bytes."""
612 if isinstance(value, str):
613 return value.encode("utf-8")
614 return value
617def make_attribute(
618 key: str,
619 value: Any,
620 doc_string: str | None = None,
621 attr_type: int | None = None,
622) -> AttributeProto:
623 """Makes an AttributeProto based on the value type."""
624 attr = AttributeProto()
625 attr.name = key
626 if doc_string:
627 attr.doc_string = doc_string
629 # Singular cases
630 if isinstance(value, numbers.Integral):
631 attr.i = int(value)
632 attr.type = AttributeProto.INT
633 elif isinstance(value, numbers.Real):
634 attr.f = float(value)
635 attr.type = AttributeProto.FLOAT
636 elif isinstance(value, (str, bytes)):
637 # Encode strings into utf-8
638 attr.s = _to_bytes(value)
639 attr.type = AttributeProto.STRING
640 elif isinstance(value, TensorProto):
641 attr.t.CopyFrom(value)
642 attr.type = AttributeProto.TENSOR
643 elif isinstance(value, onnx.SparseTensorProto):
644 attr.sparse_tensor.CopyFrom(value)
645 attr.type = AttributeProto.SPARSE_TENSOR
646 elif isinstance(value, GraphProto):
647 attr.g.CopyFrom(value)
648 attr.type = AttributeProto.GRAPH
649 elif isinstance(value, TypeProto):
650 attr.tp.CopyFrom(value)
651 attr.type = AttributeProto.TYPE_PROTO
652 # Iterable cases
653 elif isinstance(value, collections.abc.Iterable):
654 value = list(value)
655 if len(value) == 0 and attr_type is None:
656 raise ValueError(
657 f"Could not infer attribute `{key}` type from empty iterator"
658 )
659 if attr_type is None:
660 types = {type(v) for v in value}
661 for exp_t, exp_enum in (
662 (numbers.Integral, AttributeProto.INTS),
663 (numbers.Real, AttributeProto.FLOATS),
664 ((str, bytes), AttributeProto.STRINGS),
665 (TensorProto, AttributeProto.TENSORS),
666 (onnx.SparseTensorProto, AttributeProto.SPARSE_TENSORS),
667 (GraphProto, AttributeProto.GRAPHS),
668 (TypeProto, AttributeProto.TYPE_PROTOS),
669 ):
670 if all(issubclass(t, exp_t) for t in types):
671 attr_type = exp_enum
672 break
673 if attr_type is None:
674 raise ValueError(
675 "Could not infer the attribute type from the elements of the passed Iterable value."
676 )
678 if attr_type == AttributeProto.INTS:
679 # bool is a numbers.Integral, but protobuf>=7 rejects bools in int fields
680 attr.ints.extend(int(v) if isinstance(v, bool) else v for v in value)
681 attr.type = AttributeProto.INTS
682 elif attr_type == AttributeProto.FLOATS:
683 attr.floats.extend(value)
684 attr.type = AttributeProto.FLOATS
685 elif attr_type == AttributeProto.STRINGS:
686 attr.strings.extend(_to_bytes(v) for v in value)
687 attr.type = AttributeProto.STRINGS
688 elif attr_type == AttributeProto.TENSORS:
689 attr.tensors.extend(value)
690 attr.type = AttributeProto.TENSORS
691 elif attr_type == AttributeProto.SPARSE_TENSORS:
692 attr.sparse_tensors.extend(value)
693 attr.type = AttributeProto.SPARSE_TENSORS
694 elif attr_type == AttributeProto.GRAPHS:
695 attr.graphs.extend(value)
696 attr.type = AttributeProto.GRAPHS
697 elif attr_type == AttributeProto.TYPE_PROTOS:
698 attr.type_protos.extend(value)
699 attr.type = AttributeProto.TYPE_PROTOS
700 else:
701 raise AssertionError # Should not reach since `ValueError` must be raised in attr_type checking
702 else:
703 raise TypeError(f"'{value}' is not an accepted attribute value.")
705 if attr_type is not None and attr.type != attr_type:
706 raise TypeError(
707 f"Inferred attribute type '{_attr_type_to_str(attr.type)}'({attr.type}) mismatched with specified type '{_attr_type_to_str(attr_type)}'({attr_type})"
708 )
709 return attr
712def make_attribute_ref(
713 name: str,
714 attr_type: AttributeProto.AttributeType,
715 doc_string: str | None = None,
716 *,
717 ref_attr_name: str | None = None,
718) -> AttributeProto:
719 """Make an AttributeProto holding a reference to the parent function's attribute.
721 The returned attribute carries no value of its own; at instantiation time its
722 value is supplied by the parent function's attribute named ``ref_attr_name``.
723 When ``ref_attr_name`` is not provided, it defaults to ``name``. Reference
724 attributes are only valid inside a function (sub-graph).
726 Args:
727 name: The name of this attribute as used inside the function body.
728 attr_type: The type of the attribute.
729 doc_string: Optional human-readable documentation for the attribute.
730 ref_attr_name: The name of the parent function's attribute being referenced.
731 """
732 if ref_attr_name is None:
733 ref_attr_name = name
734 if not ref_attr_name:
735 raise ValueError("ref_attr_name must be non-empty")
737 attr = AttributeProto()
738 attr.name = name
739 attr.type = attr_type # type: ignore[assignment]
740 attr.ref_attr_name = ref_attr_name
741 if doc_string:
742 attr.doc_string = doc_string
743 return attr
746def get_attribute_value(attr: AttributeProto) -> Any: # noqa: PLR0911
747 if attr.ref_attr_name:
748 raise ValueError(f"Cannot get value of reference attribute: {attr}")
749 if attr.type == AttributeProto.FLOAT:
750 return attr.f
751 if attr.type == AttributeProto.INT:
752 return attr.i
753 if attr.type == AttributeProto.STRING:
754 return attr.s
755 if attr.type == AttributeProto.TENSOR:
756 return attr.t
757 if attr.type == AttributeProto.SPARSE_TENSOR:
758 return attr.sparse_tensor
759 if attr.type == AttributeProto.GRAPH:
760 return attr.g
761 if attr.type == AttributeProto.TYPE_PROTO:
762 return attr.tp
763 if attr.type == AttributeProto.FLOATS:
764 return list(attr.floats)
765 if attr.type == AttributeProto.INTS:
766 return list(attr.ints)
767 if attr.type == AttributeProto.STRINGS:
768 return list(attr.strings)
769 if attr.type == AttributeProto.TENSORS:
770 return list(attr.tensors)
771 if attr.type == AttributeProto.SPARSE_TENSORS:
772 return list(attr.sparse_tensors)
773 if attr.type == AttributeProto.GRAPHS:
774 return list(attr.graphs)
775 if attr.type == AttributeProto.TYPE_PROTOS:
776 return list(attr.type_protos)
777 if attr.type == AttributeProto.UNDEFINED:
778 return None
779 raise ValueError(f"Unsupported ONNX attribute: {attr}")
782def get_node_attr_value(node: NodeProto, attr_name: str) -> Any:
783 matching = [x for x in node.attribute if x.name == attr_name]
784 if len(matching) > 1:
785 raise ValueError(f"Node has multiple attributes with name {attr_name}")
786 if len(matching) < 1:
787 raise ValueError(f"Node has no attribute with name {attr_name}")
788 return get_attribute_value(matching[0])
791def make_empty_tensor_value_info(name: str) -> ValueInfoProto:
792 value_info_proto = ValueInfoProto()
793 value_info_proto.name = name
794 return value_info_proto
797def make_tensor_type_proto(
798 elem_type: int,
799 shape: Sequence[str | int | None] | None,
800 shape_denotation: list[str] | None = None,
801) -> TypeProto:
802 """Makes a Tensor TypeProto based on the data type and shape."""
803 type_proto = TypeProto()
804 tensor_type_proto = type_proto.tensor_type
805 tensor_type_proto.elem_type = elem_type
806 tensor_shape_proto = tensor_type_proto.shape
808 if shape is not None:
809 # You might think this is a no-op (extending a normal Python
810 # list by [] certainly is), but protobuf lists work a little
811 # differently; if a field is never set, it is omitted from the
812 # resulting protobuf; a list that is explicitly set to be
813 # empty will get an (empty) entry in the protobuf. This
814 # difference is visible to our consumers, so make sure we emit
815 # an empty shape!
816 tensor_shape_proto.dim.extend([])
818 if shape_denotation and len(shape_denotation) != len(shape):
819 raise ValueError(
820 "Invalid shape_denotation. Must be of the same length as shape."
821 )
823 for i, d in enumerate(shape):
824 dim = tensor_shape_proto.dim.add()
825 if d is None:
826 pass
827 elif isinstance(d, int):
828 dim.dim_value = d
829 elif isinstance(d, str):
830 dim.dim_param = d
831 else:
832 raise ValueError(
833 f"Invalid item in shape: {d}. Needs to be of int or str."
834 )
836 if shape_denotation:
837 dim.denotation = shape_denotation[i]
839 return type_proto
842def make_tensor_value_info(
843 name: str,
844 elem_type: int,
845 shape: Sequence[str | int | None] | None,
846 doc_string: str = "",
847 shape_denotation: list[str] | None = None,
848) -> ValueInfoProto:
849 """Makes a ValueInfoProto based on the data type and shape."""
850 value_info_proto = ValueInfoProto()
851 value_info_proto.name = name
852 if doc_string:
853 value_info_proto.doc_string = doc_string
855 tensor_type_proto = make_tensor_type_proto(elem_type, shape, shape_denotation)
856 value_info_proto.type.CopyFrom(tensor_type_proto)
857 return value_info_proto
860def make_sparse_tensor_type_proto(
861 elem_type: int,
862 shape: Sequence[str | int | None] | None,
863 shape_denotation: list[str] | None = None,
864) -> TypeProto:
865 """Makes a SparseTensor TypeProto based on the data type and shape."""
866 type_proto = TypeProto()
867 sparse_tensor_type_proto = type_proto.sparse_tensor_type
868 sparse_tensor_type_proto.elem_type = elem_type
869 sparse_tensor_shape_proto = sparse_tensor_type_proto.shape
871 if shape is not None:
872 # You might think this is a no-op (extending a normal Python
873 # list by [] certainly is), but protobuf lists work a little
874 # differently; if a field is never set, it is omitted from the
875 # resulting protobuf; a list that is explicitly set to be
876 # empty will get an (empty) entry in the protobuf. This
877 # difference is visible to our consumers, so make sure we emit
878 # an empty shape!
879 sparse_tensor_shape_proto.dim.extend([])
881 if shape_denotation and len(shape_denotation) != len(shape):
882 raise ValueError(
883 "Invalid shape_denotation. Must be of the same length as shape."
884 )
886 for i, d in enumerate(shape):
887 dim = sparse_tensor_shape_proto.dim.add()
888 if d is None:
889 pass
890 elif isinstance(d, int):
891 dim.dim_value = d
892 elif isinstance(d, str):
893 dim.dim_param = d
894 else:
895 raise ValueError(
896 f"Invalid item in shape: {d}. Needs to be of int or text."
897 )
899 if shape_denotation:
900 dim.denotation = shape_denotation[i]
902 return type_proto
905def make_sparse_tensor_value_info(
906 name: str,
907 elem_type: int,
908 shape: Sequence[str | int | None] | None,
909 doc_string: str = "",
910 shape_denotation: list[str] | None = None,
911) -> ValueInfoProto:
912 """Makes a SparseTensor ValueInfoProto based on the data type and shape."""
913 value_info_proto = ValueInfoProto()
914 value_info_proto.name = name
915 if doc_string:
916 value_info_proto.doc_string = doc_string
918 sparse_tensor_type_proto = make_sparse_tensor_type_proto(
919 elem_type, shape, shape_denotation
920 )
921 value_info_proto.type.sparse_tensor_type.CopyFrom(
922 sparse_tensor_type_proto.sparse_tensor_type
923 )
924 return value_info_proto
927def make_sequence_type_proto(
928 inner_type_proto: TypeProto,
929) -> TypeProto:
930 """Makes a sequence TypeProto."""
931 type_proto = TypeProto()
932 type_proto.sequence_type.elem_type.CopyFrom(inner_type_proto)
933 return type_proto
936def make_optional_type_proto(
937 inner_type_proto: TypeProto,
938) -> TypeProto:
939 """Makes an optional TypeProto."""
940 type_proto = TypeProto()
941 type_proto.optional_type.elem_type.CopyFrom(inner_type_proto)
942 return type_proto
945def make_map_type_proto(
946 key_type: int,
947 value_type: TypeProto,
948) -> TypeProto:
949 """Makes a map TypeProto."""
950 type_proto = TypeProto()
951 type_proto.map_type.key_type = key_type
952 type_proto.map_type.value_type.CopyFrom(value_type)
953 return type_proto
956def make_value_info(
957 name: str,
958 type_proto: TypeProto,
959 doc_string: str = "",
960) -> ValueInfoProto:
961 """Makes a ValueInfoProto with the given type_proto."""
962 value_info_proto = ValueInfoProto()
963 value_info_proto.name = name
964 if doc_string:
965 value_info_proto.doc_string = doc_string
967 value_info_proto.type.CopyFrom(type_proto)
968 return value_info_proto
971def _sanitize_str(s: str | bytes) -> str:
972 if isinstance(s, str):
973 sanitized = s
974 elif isinstance(s, bytes):
975 sanitized = s.decode("utf-8", errors="ignore")
976 else:
977 sanitized = str(s)
978 if len(sanitized) < 64: # noqa: PLR2004
979 return sanitized
980 return sanitized[:64] + f"...<+len={(len(sanitized) - 64)}>"
983def make_tensor_sequence_value_info(
984 name: str,
985 elem_type: int,
986 shape: Sequence[str | int | None] | None,
987 doc_string: str = "",
988 elem_shape_denotation: list[str] | None = None,
989) -> ValueInfoProto:
990 """Makes a Sequence[Tensors] ValueInfoProto based on the data type and shape."""
991 value_info_proto = ValueInfoProto()
992 value_info_proto.name = name
993 if doc_string:
994 value_info_proto.doc_string = doc_string
996 tensor_type_proto = make_tensor_type_proto(elem_type, shape, elem_shape_denotation)
997 sequence_type_proto = make_sequence_type_proto(tensor_type_proto)
998 value_info_proto.type.sequence_type.CopyFrom(sequence_type_proto.sequence_type)
1000 return value_info_proto
1003def printable_attribute(
1004 attr: AttributeProto, subgraphs: bool = False
1005) -> str | tuple[str, list[GraphProto]]:
1006 content = []
1007 content.append(attr.name)
1008 content.append("=")
1010 def str_float(f: float) -> str:
1011 # NB: Different Python versions print different numbers of trailing
1012 # decimals, specifying this explicitly keeps it consistent for all
1013 # versions
1014 return f"{f:.15g}"
1016 def str_int(i: int) -> str:
1017 return str(i)
1019 _T = TypeVar("_T")
1021 def str_list(str_elem: Callable[[_T], str], xs: Sequence[_T]) -> str:
1022 return "[" + ", ".join(map(str_elem, xs)) + "]"
1024 # for now, this logic should continue to work as long as we are running on a proto3
1025 # implementation. If/when we switch to proto3, we will need to use attr.type
1027 # To support printing subgraphs, if we find a graph attribute, print out
1028 # its name here and pass the graph itself up to the caller for later
1029 # printing.
1030 graphs = []
1031 if attr.HasField("f"):
1032 content.append(str_float(attr.f))
1033 elif attr.HasField("i"):
1034 content.append(str_int(attr.i))
1035 elif attr.HasField("s"):
1036 content.append(repr(_sanitize_str(attr.s)))
1037 elif attr.HasField("t"):
1038 if len(attr.t.dims) > 0:
1039 content.append("<Tensor>")
1040 else:
1041 # special case to print scalars
1042 field = tensor_dtype_to_field(attr.t.data_type)
1043 content.append(f"<Scalar Tensor {getattr(attr.t, field)}>")
1044 elif attr.HasField("g"):
1045 content.append(f"<graph {attr.g.name}>")
1046 graphs.append(attr.g)
1047 elif attr.HasField("tp"):
1048 content.append(f"<Type Proto {attr.tp}>")
1049 elif attr.HasField("sparse_tensor"):
1050 content.append("<Sparse Tensor>")
1051 elif attr.floats:
1052 content.append(str_list(str_float, attr.floats))
1053 elif attr.ints:
1054 content.append(str_list(str_int, attr.ints))
1055 elif attr.strings:
1056 content.append(str(list(map(_sanitize_str, attr.strings))))
1057 elif attr.tensors:
1058 content.append("[<Tensor>, ...]")
1059 elif attr.sparse_tensors:
1060 content.append("[<Sparse Tensor>, ...]")
1061 elif attr.type_protos:
1062 content.append("[")
1063 for i, tp in enumerate(attr.type_protos):
1064 comma = "," if i != len(attr.type_protos) - 1 else ""
1065 content.append(f"<Type Proto {tp}>{comma}")
1066 content.append("]")
1067 elif attr.graphs:
1068 content.append("[")
1069 for i, g in enumerate(attr.graphs):
1070 comma = "," if i != len(attr.graphs) - 1 else ""
1071 content.append(f"<graph {g.name}>{comma}")
1072 content.append("]")
1073 graphs.extend(attr.graphs)
1074 else:
1075 content.append("<Unknown>")
1076 if subgraphs:
1077 return " ".join(content), graphs
1078 return " ".join(content)
1081def printable_dim(dim: TensorShapeProto.Dimension) -> str:
1082 which = dim.WhichOneof("value")
1083 if which is None:
1084 return "?"
1085 return str(getattr(dim, which))
1088def printable_type(t: TypeProto) -> str:
1089 if t.WhichOneof("value") == "tensor_type":
1090 s: str = TensorProto.DataType.Name(t.tensor_type.elem_type) # type: ignore[arg-type]
1091 if t.tensor_type.HasField("shape"):
1092 if len(t.tensor_type.shape.dim):
1093 s += str(", " + "x".join(map(printable_dim, t.tensor_type.shape.dim)))
1094 else:
1095 s += ", scalar"
1096 return s
1097 if t.WhichOneof("value") is None:
1098 return ""
1099 return f"Unknown type {t.WhichOneof('value')}"
1102def printable_value_info(v: ValueInfoProto) -> str:
1103 s = f"%{v.name}"
1104 if v.type:
1105 s = f"{s}[{printable_type(v.type)}]"
1106 return s
1109def printable_tensor_proto(t: TensorProto) -> str:
1110 s = f"%{t.name}["
1111 s += TensorProto.DataType.Name(t.data_type) # type: ignore[arg-type]
1112 if t.dims is not None:
1113 if len(t.dims):
1114 s += str(", " + "x".join(map(str, t.dims)))
1115 else:
1116 s += ", scalar"
1117 s += "]"
1118 return s
1121def printable_node(
1122 node: NodeProto, prefix: str = "", subgraphs: bool = False
1123) -> str | tuple[str, list[GraphProto]]:
1124 content = []
1125 if len(node.output):
1126 content.append(", ".join([f"%{name}" for name in node.output]))
1127 content.append("=")
1128 # To deal with nested graphs
1129 graphs: list[GraphProto] = []
1130 printed_attrs = []
1131 for attr in node.attribute:
1132 if subgraphs:
1133 printed_attr_subgraphs = printable_attribute(attr, subgraphs)
1134 if not isinstance(printed_attr_subgraphs[1], list):
1135 raise TypeError(
1136 f"printed_attr_subgraphs[1] must be an instance of {list}."
1137 )
1138 graphs.extend(printed_attr_subgraphs[1])
1139 printed_attrs.append(printed_attr_subgraphs[0])
1140 else:
1141 printed = printable_attribute(attr)
1142 if not isinstance(printed, str):
1143 raise TypeError(f"printed must be an instance of {str}.")
1144 printed_attrs.append(printed)
1145 printed_attributes = ", ".join(sorted(printed_attrs))
1146 printed_inputs = ", ".join([f"%{name}" for name in node.input])
1147 if node.attribute:
1148 content.append(f"{node.op_type}[{printed_attributes}]({printed_inputs})")
1149 else:
1150 content.append(f"{node.op_type}({printed_inputs})")
1151 if subgraphs:
1152 return prefix + " ".join(content), graphs
1153 return prefix + " ".join(content)
1156@typing_extensions.deprecated(
1157 "Deprecated since 1.19. Consider using onnx.printer.to_text() instead."
1158)
1159def printable_graph(graph: GraphProto, prefix: str = "") -> str:
1160 """Display a GraphProto as a string.
1162 .. deprecated:: 1.19
1163 Consider using :func:`onnx.printer.to_text` instead.
1165 Args:
1166 graph (GraphProto): the graph to display
1167 prefix (string): prefix of every line
1169 Returns:
1170 string
1171 """
1172 content = []
1173 indent = prefix + " "
1174 # header
1175 header = ["graph", graph.name]
1176 initializers = {t.name for t in graph.initializer}
1177 if len(graph.input):
1178 header.append("(")
1179 in_strs = [] # required inputs
1180 in_with_init_strs: list = [] # optional inputs with initializer providing default value
1181 for inp in graph.input:
1182 if inp.name not in initializers:
1183 in_strs.append(printable_value_info(inp))
1184 else:
1185 in_with_init_strs.append(printable_value_info(inp))
1186 if in_strs:
1187 content.append(prefix + " ".join(header))
1188 header = []
1189 for line in in_strs:
1190 content.append(prefix + " " + line) # noqa: PERF401
1191 header.append(")")
1193 if in_with_init_strs:
1194 header.append("optional inputs with matching initializers (")
1195 content.append(prefix + " ".join(header))
1196 header = []
1197 for line in in_with_init_strs:
1198 content.append(prefix + " " + line) # noqa: PERF401
1199 header.append(")")
1201 # from IR 4 onwards an initializer is not required to have a matching graph input
1202 # so output the name, type and shape of those as well
1203 if len(in_with_init_strs) < len(initializers):
1204 graph_inputs = {i.name for i in graph.input}
1205 init_strs = [
1206 printable_tensor_proto(i)
1207 for i in graph.initializer
1208 if i.name not in graph_inputs
1209 ]
1210 header.append("initializers (")
1211 content.append(prefix + " ".join(header))
1212 header = []
1213 for line in init_strs:
1214 content.append(prefix + " " + line) # noqa: PERF401
1215 header.append(")")
1217 header.append("{")
1218 content.append(prefix + " ".join(header))
1219 graphs: list[GraphProto] = []
1220 # body
1221 for node in graph.node:
1222 contents_subgraphs = printable_node(node, indent, subgraphs=True)
1223 if not isinstance(contents_subgraphs[1], list):
1224 raise TypeError(f"contents_subgraphs[1] must be an instance of {list}.")
1225 content.append(contents_subgraphs[0])
1226 graphs.extend(contents_subgraphs[1])
1227 # tail
1228 tail = ["return"]
1229 if len(graph.output):
1230 tail.append(", ".join([f"%{out.name}" for out in graph.output]))
1231 content.append(indent + " ".join(tail))
1232 # closing bracket
1233 content.append(prefix + "}")
1234 for g in graphs:
1235 content.append("\n" + printable_graph(g)) # noqa: PERF401
1236 return "\n".join(content)
1239def strip_doc_string(proto: google.protobuf.message.Message) -> None:
1240 """Empties `doc_string` field on any nested protobuf messages"""
1241 if not isinstance(proto, google.protobuf.message.Message):
1242 raise TypeError(
1243 f"proto must be an instance of {google.protobuf.message.Message}."
1244 )
1245 for descriptor in proto.DESCRIPTOR.fields:
1246 if descriptor.name == "doc_string":
1247 proto.ClearField(descriptor.name)
1248 elif descriptor.type == descriptor.TYPE_MESSAGE:
1249 if descriptor.is_repeated:
1250 for x in getattr(proto, descriptor.name):
1251 strip_doc_string(x)
1252 elif proto.HasField(descriptor.name):
1253 strip_doc_string(getattr(proto, descriptor.name))
1256def make_training_info(
1257 algorithm: GraphProto,
1258 algorithm_bindings: AssignmentBindingType,
1259 initialization: GraphProto | None,
1260 initialization_bindings: AssignmentBindingType | None,
1261) -> TrainingInfoProto:
1262 training_info = TrainingInfoProto()
1263 training_info.algorithm.CopyFrom(algorithm)
1264 for k, v in algorithm_bindings:
1265 binding = training_info.update_binding.add()
1266 binding.key = k
1267 binding.value = v
1269 if initialization:
1270 training_info.initialization.CopyFrom(initialization)
1271 if initialization_bindings:
1272 for k, v in initialization_bindings:
1273 binding = training_info.initialization_binding.add()
1274 binding.key = k
1275 binding.value = v
1277 return training_info
1280# Following functions are used for mapping
1281def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
1282 """Convert a TensorProto's data_type to corresponding numpy dtype. It can be used while making tensor.
1284 Args:
1285 tensor_dtype: TensorProto's data_type
1287 Returns:
1288 numpy's data_type
1289 """
1290 return _mapping.TENSOR_TYPE_MAP[tensor_dtype].np_dtype
1293def tensor_dtype_to_storage_tensor_dtype(tensor_dtype: int) -> int:
1294 """Convert a TensorProto's data_type to corresponding data_type for storage.
1296 Args:
1297 tensor_dtype: TensorProto's data_type
1299 Returns:
1300 data_type for storage
1301 """
1302 return _mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
1305def tensor_dtype_to_string(tensor_dtype: int) -> str:
1306 """Get the name of given TensorProto's data_type.
1308 Args:
1309 tensor_dtype: TensorProto's data_type
1311 Returns:
1312 the name of data_type
1313 """
1314 return _mapping.TENSOR_TYPE_MAP[tensor_dtype].name
1317@functools.lru_cache(None)
1318def tensor_dtype_to_field(tensor_dtype: int) -> str:
1319 """Convert a TensorProto's data_type to corresponding field name for storage. It can be used while making tensors.
1321 Args:
1322 tensor_dtype: TensorProto's data_type
1324 Returns:
1325 field name
1326 """
1327 storage_tensor_type_to_field = {
1328 int(TensorProto.FLOAT): "float_data",
1329 int(TensorProto.INT32): "int32_data",
1330 int(TensorProto.INT64): "int64_data",
1331 int(TensorProto.DOUBLE): "double_data",
1332 int(TensorProto.UINT32): "uint64_data",
1333 int(TensorProto.UINT64): "uint64_data",
1334 int(TensorProto.STRING): "string_data",
1335 }
1336 return storage_tensor_type_to_field[
1337 _mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
1338 ]
1341@functools.lru_cache(None)
1342def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> TensorProto.DataType:
1343 """Convert a numpy's dtype to corresponding tensor type. It can be used while converting numpy arrays to tensors.
1345 Args:
1346 np_dtype: numpy's data_type
1348 Returns:
1349 TensorsProto's data_type
1350 """
1351 _np_dtype_to_tensor_dtype = {
1352 v.np_dtype: k for k, v in _mapping.TENSOR_TYPE_MAP.items()
1353 }
1354 if np_dtype in _np_dtype_to_tensor_dtype:
1355 return typing.cast("TensorProto.DataType", _np_dtype_to_tensor_dtype[np_dtype])
1356 if np.issubdtype(np_dtype, np.str_):
1357 return TensorProto.STRING # type: ignore[return-value]
1359 raise ValueError(
1360 f"Unable to convert type {np_dtype!r} into TensorProto element type."
1361 )
1364def get_all_tensor_dtypes() -> KeysView[int]:
1365 """Get all tensor types from TensorProto.
1367 Returns:
1368 all tensor types from TensorProto
1369 """
1370 return _mapping.TENSOR_TYPE_MAP.keys()
1373_ATTRIBUTE_TYPE_TO_STR: dict[int, str] = {
1374 k: v for v, k in AttributeProto.AttributeType.items()
1375}
1378def _attr_type_to_str(attr_type: int) -> str:
1379 """Convert AttributeProto type to string.
1381 Args:
1382 attr_type: AttributeProto type.
1384 Returns:
1385 String representing the supplied attr_type.
1386 """
1387 if attr_type in AttributeProto.AttributeType.values():
1388 return _ATTRIBUTE_TYPE_TO_STR[attr_type]
1389 return AttributeProto.AttributeType.keys()[0]