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", 13, 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. It is packed in couples to UINT8 buffers.
401 if data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
402 expected_size_bytes = 0.5
403 elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
404 expected_size_bytes = 0.25
405 else:
406 expected_size_bytes = np_dtype.itemsize
407 expected_size_bytes *= math.prod(dims)
408 expected_size_bytes = math.ceil(expected_size_bytes)
409 if isinstance(vals, np.ndarray):
410 if data_type in {
411 TensorProto.INT4,
412 TensorProto.UINT4,
413 TensorProto.FLOAT4E2M1,
414 }:
415 vals = onnx.numpy_helper._pack_4bitx2(vals)
416 elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
417 vals = onnx.numpy_helper._pack_2bitx4(vals)
419 raw_data = onnx.numpy_helper.tobytes_little_endian(vals)
420 elif isinstance(vals, bytes):
421 raw_data = vals
422 else:
423 raise TypeError(
424 f"Raw data must be bytes or numpy.ndarray, but got {type(vals)}."
425 )
426 if len(raw_data) != expected_size_bytes:
427 raise ValueError(
428 f"Raw data size does not match tensor's size. Expected {expected_size_bytes} bytes, but got {len(raw_data)} bytes."
429 )
430 tensor.raw_data = raw_data
431 return tensor
433 assert not raw, "Bug: raw should be False at this point."
435 if data_type == TensorProto.STRING:
436 vals = np.array(vals).flatten()
437 if len(vals) != 0:
438 vals = np.vectorize(_to_bytes)(vals) # Convert to bytes
439 elif data_type in {
440 TensorProto.FLOAT8E4M3FN,
441 TensorProto.FLOAT8E4M3FNUZ,
442 TensorProto.FLOAT8E5M2,
443 TensorProto.FLOAT8E5M2FNUZ,
444 }:
445 # Float8 values are by default casted using saturating cast.
446 vals = onnx.numpy_helper.saturate_cast(np.asarray(vals), np_dtype).flatten()
447 elif data_type == TensorProto.FLOAT8E8M0:
448 vals = onnx.numpy_helper.to_float8e8m0(
449 np.asarray(vals), saturate=True, round_mode="up"
450 ).flatten()
451 else:
452 vals = np.asarray(vals, dtype=np_dtype).flatten()
454 expected_elements = math.prod(dims)
455 if len(vals) != expected_elements:
456 raise ValueError(
457 f"Number of values ({len(vals)}) does not match tensor "
458 f"dimensions requiring {expected_elements} elements."
459 )
460 if data_type == TensorProto.COMPLEX128:
461 vals = vals.view(np.float64) # type: ignore[union-attr]
462 elif data_type == TensorProto.COMPLEX64:
463 vals = vals.view(np.float32) # type: ignore[union-attr]
464 elif data_type in {TensorProto.BFLOAT16, TensorProto.FLOAT16}:
465 vals = vals.view(np.uint16) # type: ignore[union-attr]
466 elif data_type in {
467 TensorProto.FLOAT8E4M3FN,
468 TensorProto.FLOAT8E4M3FNUZ,
469 TensorProto.FLOAT8E5M2,
470 TensorProto.FLOAT8E5M2FNUZ,
471 TensorProto.FLOAT8E8M0,
472 }:
473 vals = vals.view(np.uint8) # type: ignore[union-attr]
474 elif data_type in {TensorProto.UINT4, TensorProto.INT4, TensorProto.FLOAT4E2M1}:
475 # Convert to packed 4-bit representation
476 vals = onnx.numpy_helper._pack_4bitx2(vals) # type: ignore[arg-type]
477 elif data_type in {TensorProto.UINT2, TensorProto.INT2}:
478 # Convert to packed 2-bit representation
479 vals = onnx.numpy_helper._pack_2bitx4(vals) # type: ignore[arg-type]
480 elif data_type == TensorProto.BOOL:
481 vals = vals.astype(np.uint8) # type: ignore[union-attr]
483 field = tensor_dtype_to_field(data_type)
484 getattr(tensor, field).extend(vals)
485 return tensor
488def make_sparse_tensor(
489 values: TensorProto, indices: TensorProto, dims: Sequence[int]
490) -> onnx.SparseTensorProto:
491 """Construct a SparseTensorProto
493 Args:
494 values (TensorProto): the values
495 indices (TensorProto): the indices
496 dims: the shape
498 Returns:
499 SparseTensorProto
500 """
501 sparse = onnx.SparseTensorProto()
502 sparse.values.CopyFrom(values)
503 sparse.indices.CopyFrom(indices)
504 sparse.dims.extend(dims)
505 return sparse
508def make_sequence(
509 name: str,
510 elem_type: SequenceProto.DataType,
511 values: Sequence[Any],
512) -> SequenceProto:
513 """Make a Sequence with specified value arguments."""
514 sequence = SequenceProto()
515 sequence.name = name
516 sequence.elem_type = elem_type # type: ignore[assignment]
518 if elem_type == SequenceProto.UNDEFINED:
519 return sequence
521 attribute: RepeatedCompositeFieldContainer | None = None
522 if elem_type == SequenceProto.TENSOR:
523 attribute = sequence.tensor_values
524 elif elem_type == SequenceProto.SPARSE_TENSOR:
525 attribute = sequence.sparse_tensor_values
526 elif elem_type == SequenceProto.SEQUENCE:
527 attribute = sequence.sequence_values
528 elif elem_type == SequenceProto.MAP:
529 attribute = sequence.map_values
530 elif elem_type == OptionalProto.OPTIONAL:
531 attribute = sequence.optional_values
532 else:
533 raise TypeError("The element type in the input sequence is not supported.")
535 attribute.extend(values)
536 return sequence
539def make_map(
540 name: str, key_type: int, keys: list[Any], values: SequenceProto
541) -> MapProto:
542 """Make a Map with specified key-value pair arguments.
544 Criteria for conversion:
545 - Keys and Values must have the same number of elements
546 - Every key in keys must be of the same type
547 - Every value in values must be of the same type
548 """
549 map_proto = MapProto()
550 valid_key_int_types = [
551 TensorProto.INT8,
552 TensorProto.INT16,
553 TensorProto.INT32,
554 TensorProto.INT64,
555 TensorProto.UINT8,
556 TensorProto.UINT16,
557 TensorProto.UINT32,
558 TensorProto.UINT64,
559 ]
560 map_proto.name = name
561 map_proto.key_type = key_type
562 if key_type == TensorProto.STRING:
563 map_proto.string_keys.extend(keys)
564 elif key_type in valid_key_int_types:
565 map_proto.keys.extend(keys)
566 map_proto.values.CopyFrom(values)
567 return map_proto
570def make_optional(
571 name: str,
572 elem_type: OptionalProto.DataType,
573 value: google.protobuf.message.Message | None,
574) -> OptionalProto:
575 """Make an Optional with specified value arguments."""
576 optional = OptionalProto()
577 optional.name = name
578 optional.elem_type = elem_type # type: ignore[assignment]
580 if elem_type == OptionalProto.UNDEFINED:
581 return optional
582 attribute: google.protobuf.message.Message | None = None
583 if elem_type == OptionalProto.TENSOR:
584 attribute = optional.tensor_value
585 elif elem_type == OptionalProto.SPARSE_TENSOR:
586 attribute = optional.sparse_tensor_value
587 elif elem_type == OptionalProto.SEQUENCE:
588 attribute = optional.sequence_value
589 elif elem_type == OptionalProto.MAP:
590 attribute = optional.map_value
591 elif elem_type == OptionalProto.OPTIONAL:
592 attribute = optional.optional_value
593 else:
594 raise TypeError("The element type in the input optional is not supported.")
596 assert value is not None
597 attribute.CopyFrom(value) # type: ignore[arg-type]
598 return optional
601def _to_bytes(value: str | bytes) -> bytes:
602 """Coerce a string (or bytes) value into UTF-8 bytes."""
603 if isinstance(value, str):
604 return value.encode("utf-8")
605 return value
608def make_attribute(
609 key: str,
610 value: Any,
611 doc_string: str | None = None,
612 attr_type: int | None = None,
613) -> AttributeProto:
614 """Makes an AttributeProto based on the value type."""
615 attr = AttributeProto()
616 attr.name = key
617 if doc_string:
618 attr.doc_string = doc_string
620 # Singular cases
621 if isinstance(value, numbers.Integral):
622 attr.i = int(value)
623 attr.type = AttributeProto.INT
624 elif isinstance(value, numbers.Real):
625 attr.f = float(value)
626 attr.type = AttributeProto.FLOAT
627 elif isinstance(value, (str, bytes)):
628 # Encode strings into utf-8
629 attr.s = _to_bytes(value)
630 attr.type = AttributeProto.STRING
631 elif isinstance(value, TensorProto):
632 attr.t.CopyFrom(value)
633 attr.type = AttributeProto.TENSOR
634 elif isinstance(value, onnx.SparseTensorProto):
635 attr.sparse_tensor.CopyFrom(value)
636 attr.type = AttributeProto.SPARSE_TENSOR
637 elif isinstance(value, GraphProto):
638 attr.g.CopyFrom(value)
639 attr.type = AttributeProto.GRAPH
640 elif isinstance(value, TypeProto):
641 attr.tp.CopyFrom(value)
642 attr.type = AttributeProto.TYPE_PROTO
643 # Iterable cases
644 elif isinstance(value, collections.abc.Iterable):
645 value = list(value)
646 if len(value) == 0 and attr_type is None:
647 raise ValueError(
648 f"Could not infer attribute `{key}` type from empty iterator"
649 )
650 if attr_type is None:
651 types = {type(v) for v in value}
652 for exp_t, exp_enum in (
653 (numbers.Integral, AttributeProto.INTS),
654 (numbers.Real, AttributeProto.FLOATS),
655 ((str, bytes), AttributeProto.STRINGS),
656 (TensorProto, AttributeProto.TENSORS),
657 (onnx.SparseTensorProto, AttributeProto.SPARSE_TENSORS),
658 (GraphProto, AttributeProto.GRAPHS),
659 (TypeProto, AttributeProto.TYPE_PROTOS),
660 ):
661 if all(issubclass(t, exp_t) for t in types):
662 attr_type = exp_enum
663 break
664 if attr_type is None:
665 raise ValueError(
666 "Could not infer the attribute type from the elements of the passed Iterable value."
667 )
669 if attr_type == AttributeProto.INTS:
670 attr.ints.extend(value)
671 attr.type = AttributeProto.INTS
672 elif attr_type == AttributeProto.FLOATS:
673 attr.floats.extend(value)
674 attr.type = AttributeProto.FLOATS
675 elif attr_type == AttributeProto.STRINGS:
676 attr.strings.extend(_to_bytes(v) for v in value)
677 attr.type = AttributeProto.STRINGS
678 elif attr_type == AttributeProto.TENSORS:
679 attr.tensors.extend(value)
680 attr.type = AttributeProto.TENSORS
681 elif attr_type == AttributeProto.SPARSE_TENSORS:
682 attr.sparse_tensors.extend(value)
683 attr.type = AttributeProto.SPARSE_TENSORS
684 elif attr_type == AttributeProto.GRAPHS:
685 attr.graphs.extend(value)
686 attr.type = AttributeProto.GRAPHS
687 elif attr_type == AttributeProto.TYPE_PROTOS:
688 attr.type_protos.extend(value)
689 attr.type = AttributeProto.TYPE_PROTOS
690 else:
691 raise AssertionError # Should not reach since `ValueError` must be raised in attr_type checking
692 else:
693 raise TypeError(f"'{value}' is not an accepted attribute value.")
695 if attr_type is not None and attr.type != attr_type:
696 raise TypeError(
697 f"Inferred attribute type '{_attr_type_to_str(attr.type)}'({attr.type}) mismatched with specified type '{_attr_type_to_str(attr_type)}'({attr_type})"
698 )
699 return attr
702def make_attribute_ref(
703 name: str,
704 attr_type: AttributeProto.AttributeType,
705 doc_string: str | None = None,
706 *,
707 ref_attr_name: str | None = None,
708) -> AttributeProto:
709 """Make an AttributeProto holding a reference to the parent function's attribute.
711 The returned attribute carries no value of its own; at instantiation time its
712 value is supplied by the parent function's attribute named ``ref_attr_name``.
713 When ``ref_attr_name`` is not provided, it defaults to ``name``. Reference
714 attributes are only valid inside a function (sub-graph).
716 Args:
717 name: The name of this attribute as used inside the function body.
718 attr_type: The type of the attribute.
719 doc_string: Optional human-readable documentation for the attribute.
720 ref_attr_name: The name of the parent function's attribute being referenced.
721 """
722 if ref_attr_name is None:
723 ref_attr_name = name
724 if not ref_attr_name:
725 raise ValueError("ref_attr_name must be non-empty")
727 attr = AttributeProto()
728 attr.name = name
729 attr.type = attr_type # type: ignore[assignment]
730 attr.ref_attr_name = ref_attr_name
731 if doc_string:
732 attr.doc_string = doc_string
733 return attr
736def get_attribute_value(attr: AttributeProto) -> Any: # noqa: PLR0911
737 if attr.ref_attr_name:
738 raise ValueError(f"Cannot get value of reference attribute: {attr}")
739 if attr.type == AttributeProto.FLOAT:
740 return attr.f
741 if attr.type == AttributeProto.INT:
742 return attr.i
743 if attr.type == AttributeProto.STRING:
744 return attr.s
745 if attr.type == AttributeProto.TENSOR:
746 return attr.t
747 if attr.type == AttributeProto.SPARSE_TENSOR:
748 return attr.sparse_tensor
749 if attr.type == AttributeProto.GRAPH:
750 return attr.g
751 if attr.type == AttributeProto.TYPE_PROTO:
752 return attr.tp
753 if attr.type == AttributeProto.FLOATS:
754 return list(attr.floats)
755 if attr.type == AttributeProto.INTS:
756 return list(attr.ints)
757 if attr.type == AttributeProto.STRINGS:
758 return list(attr.strings)
759 if attr.type == AttributeProto.TENSORS:
760 return list(attr.tensors)
761 if attr.type == AttributeProto.SPARSE_TENSORS:
762 return list(attr.sparse_tensors)
763 if attr.type == AttributeProto.GRAPHS:
764 return list(attr.graphs)
765 if attr.type == AttributeProto.TYPE_PROTOS:
766 return list(attr.type_protos)
767 if attr.type == AttributeProto.UNDEFINED:
768 return None
769 raise ValueError(f"Unsupported ONNX attribute: {attr}")
772def get_node_attr_value(node: NodeProto, attr_name: str) -> Any:
773 matching = [x for x in node.attribute if x.name == attr_name]
774 if len(matching) > 1:
775 raise ValueError(f"Node has multiple attributes with name {attr_name}")
776 if len(matching) < 1:
777 raise ValueError(f"Node has no attribute with name {attr_name}")
778 return get_attribute_value(matching[0])
781def make_empty_tensor_value_info(name: str) -> ValueInfoProto:
782 value_info_proto = ValueInfoProto()
783 value_info_proto.name = name
784 return value_info_proto
787def make_tensor_type_proto(
788 elem_type: int,
789 shape: Sequence[str | int | None] | None,
790 shape_denotation: list[str] | None = None,
791) -> TypeProto:
792 """Makes a Tensor TypeProto based on the data type and shape."""
793 type_proto = TypeProto()
794 tensor_type_proto = type_proto.tensor_type
795 tensor_type_proto.elem_type = elem_type
796 tensor_shape_proto = tensor_type_proto.shape
798 if shape is not None:
799 # You might think this is a no-op (extending a normal Python
800 # list by [] certainly is), but protobuf lists work a little
801 # differently; if a field is never set, it is omitted from the
802 # resulting protobuf; a list that is explicitly set to be
803 # empty will get an (empty) entry in the protobuf. This
804 # difference is visible to our consumers, so make sure we emit
805 # an empty shape!
806 tensor_shape_proto.dim.extend([])
808 if shape_denotation and len(shape_denotation) != len(shape):
809 raise ValueError(
810 "Invalid shape_denotation. Must be of the same length as shape."
811 )
813 for i, d in enumerate(shape):
814 dim = tensor_shape_proto.dim.add()
815 if d is None:
816 pass
817 elif isinstance(d, int):
818 dim.dim_value = d
819 elif isinstance(d, str):
820 dim.dim_param = d
821 else:
822 raise ValueError(
823 f"Invalid item in shape: {d}. Needs to be of int or str."
824 )
826 if shape_denotation:
827 dim.denotation = shape_denotation[i]
829 return type_proto
832def make_tensor_value_info(
833 name: str,
834 elem_type: int,
835 shape: Sequence[str | int | None] | None,
836 doc_string: str = "",
837 shape_denotation: list[str] | None = None,
838) -> ValueInfoProto:
839 """Makes a ValueInfoProto based on the data type and shape."""
840 value_info_proto = ValueInfoProto()
841 value_info_proto.name = name
842 if doc_string:
843 value_info_proto.doc_string = doc_string
845 tensor_type_proto = make_tensor_type_proto(elem_type, shape, shape_denotation)
846 value_info_proto.type.CopyFrom(tensor_type_proto)
847 return value_info_proto
850def make_sparse_tensor_type_proto(
851 elem_type: int,
852 shape: Sequence[str | int | None] | None,
853 shape_denotation: list[str] | None = None,
854) -> TypeProto:
855 """Makes a SparseTensor TypeProto based on the data type and shape."""
856 type_proto = TypeProto()
857 sparse_tensor_type_proto = type_proto.sparse_tensor_type
858 sparse_tensor_type_proto.elem_type = elem_type
859 sparse_tensor_shape_proto = sparse_tensor_type_proto.shape
861 if shape is not None:
862 # You might think this is a no-op (extending a normal Python
863 # list by [] certainly is), but protobuf lists work a little
864 # differently; if a field is never set, it is omitted from the
865 # resulting protobuf; a list that is explicitly set to be
866 # empty will get an (empty) entry in the protobuf. This
867 # difference is visible to our consumers, so make sure we emit
868 # an empty shape!
869 sparse_tensor_shape_proto.dim.extend([])
871 if shape_denotation and len(shape_denotation) != len(shape):
872 raise ValueError(
873 "Invalid shape_denotation. Must be of the same length as shape."
874 )
876 for i, d in enumerate(shape):
877 dim = sparse_tensor_shape_proto.dim.add()
878 if d is None:
879 pass
880 elif isinstance(d, int):
881 dim.dim_value = d
882 elif isinstance(d, str):
883 dim.dim_param = d
884 else:
885 raise ValueError(
886 f"Invalid item in shape: {d}. Needs to be of int or text."
887 )
889 if shape_denotation:
890 dim.denotation = shape_denotation[i]
892 return type_proto
895def make_sparse_tensor_value_info(
896 name: str,
897 elem_type: int,
898 shape: Sequence[str | int | None] | None,
899 doc_string: str = "",
900 shape_denotation: list[str] | None = None,
901) -> ValueInfoProto:
902 """Makes a SparseTensor ValueInfoProto based on the data type and shape."""
903 value_info_proto = ValueInfoProto()
904 value_info_proto.name = name
905 if doc_string:
906 value_info_proto.doc_string = doc_string
908 sparse_tensor_type_proto = make_sparse_tensor_type_proto(
909 elem_type, shape, shape_denotation
910 )
911 value_info_proto.type.sparse_tensor_type.CopyFrom(
912 sparse_tensor_type_proto.sparse_tensor_type
913 )
914 return value_info_proto
917def make_sequence_type_proto(
918 inner_type_proto: TypeProto,
919) -> TypeProto:
920 """Makes a sequence TypeProto."""
921 type_proto = TypeProto()
922 type_proto.sequence_type.elem_type.CopyFrom(inner_type_proto)
923 return type_proto
926def make_optional_type_proto(
927 inner_type_proto: TypeProto,
928) -> TypeProto:
929 """Makes an optional TypeProto."""
930 type_proto = TypeProto()
931 type_proto.optional_type.elem_type.CopyFrom(inner_type_proto)
932 return type_proto
935def make_map_type_proto(
936 key_type: int,
937 value_type: TypeProto,
938) -> TypeProto:
939 """Makes a map TypeProto."""
940 type_proto = TypeProto()
941 type_proto.map_type.key_type = key_type
942 type_proto.map_type.value_type.CopyFrom(value_type)
943 return type_proto
946def make_value_info(
947 name: str,
948 type_proto: TypeProto,
949 doc_string: str = "",
950) -> ValueInfoProto:
951 """Makes a ValueInfoProto with the given type_proto."""
952 value_info_proto = ValueInfoProto()
953 value_info_proto.name = name
954 if doc_string:
955 value_info_proto.doc_string = doc_string
957 value_info_proto.type.CopyFrom(type_proto)
958 return value_info_proto
961def _sanitize_str(s: str | bytes) -> str:
962 if isinstance(s, str):
963 sanitized = s
964 elif isinstance(s, bytes):
965 sanitized = s.decode("utf-8", errors="ignore")
966 else:
967 sanitized = str(s)
968 if len(sanitized) < 64: # noqa: PLR2004
969 return sanitized
970 return sanitized[:64] + f"...<+len={(len(sanitized) - 64)}>"
973def make_tensor_sequence_value_info(
974 name: str,
975 elem_type: int,
976 shape: Sequence[str | int | None] | None,
977 doc_string: str = "",
978 elem_shape_denotation: list[str] | None = None,
979) -> ValueInfoProto:
980 """Makes a Sequence[Tensors] ValueInfoProto based on the data type and shape."""
981 value_info_proto = ValueInfoProto()
982 value_info_proto.name = name
983 if doc_string:
984 value_info_proto.doc_string = doc_string
986 tensor_type_proto = make_tensor_type_proto(elem_type, shape, elem_shape_denotation)
987 sequence_type_proto = make_sequence_type_proto(tensor_type_proto)
988 value_info_proto.type.sequence_type.CopyFrom(sequence_type_proto.sequence_type)
990 return value_info_proto
993def printable_attribute(
994 attr: AttributeProto, subgraphs: bool = False
995) -> str | tuple[str, list[GraphProto]]:
996 content = []
997 content.append(attr.name)
998 content.append("=")
1000 def str_float(f: float) -> str:
1001 # NB: Different Python versions print different numbers of trailing
1002 # decimals, specifying this explicitly keeps it consistent for all
1003 # versions
1004 return f"{f:.15g}"
1006 def str_int(i: int) -> str:
1007 return str(i)
1009 _T = TypeVar("_T")
1011 def str_list(str_elem: Callable[[_T], str], xs: Sequence[_T]) -> str:
1012 return "[" + ", ".join(map(str_elem, xs)) + "]"
1014 # for now, this logic should continue to work as long as we are running on a proto3
1015 # implementation. If/when we switch to proto3, we will need to use attr.type
1017 # To support printing subgraphs, if we find a graph attribute, print out
1018 # its name here and pass the graph itself up to the caller for later
1019 # printing.
1020 graphs = []
1021 if attr.HasField("f"):
1022 content.append(str_float(attr.f))
1023 elif attr.HasField("i"):
1024 content.append(str_int(attr.i))
1025 elif attr.HasField("s"):
1026 # TODO: Bit nervous about Python 2 / Python 3 determinism implications
1027 content.append(repr(_sanitize_str(attr.s)))
1028 elif attr.HasField("t"):
1029 if len(attr.t.dims) > 0:
1030 content.append("<Tensor>")
1031 else:
1032 # special case to print scalars
1033 field = tensor_dtype_to_field(attr.t.data_type)
1034 content.append(f"<Scalar Tensor {getattr(attr.t, field)}>")
1035 elif attr.HasField("g"):
1036 content.append(f"<graph {attr.g.name}>")
1037 graphs.append(attr.g)
1038 elif attr.HasField("tp"):
1039 content.append(f"<Type Proto {attr.tp}>")
1040 elif attr.HasField("sparse_tensor"):
1041 content.append("<Sparse Tensor>")
1042 elif attr.floats:
1043 content.append(str_list(str_float, attr.floats))
1044 elif attr.ints:
1045 content.append(str_list(str_int, attr.ints))
1046 elif attr.strings:
1047 # TODO: Bit nervous about Python 2 / Python 3 determinism implications
1048 content.append(str(list(map(_sanitize_str, attr.strings))))
1049 elif attr.tensors:
1050 content.append("[<Tensor>, ...]")
1051 elif attr.sparse_tensors:
1052 content.append("[<Sparse Tensor>, ...]")
1053 elif attr.type_protos:
1054 content.append("[")
1055 for i, tp in enumerate(attr.type_protos):
1056 comma = "," if i != len(attr.type_protos) - 1 else ""
1057 content.append(f"<Type Proto {tp}>{comma}")
1058 content.append("]")
1059 elif attr.graphs:
1060 content.append("[")
1061 for i, g in enumerate(attr.graphs):
1062 comma = "," if i != len(attr.graphs) - 1 else ""
1063 content.append(f"<graph {g.name}>{comma}")
1064 content.append("]")
1065 graphs.extend(attr.graphs)
1066 else:
1067 content.append("<Unknown>")
1068 if subgraphs:
1069 return " ".join(content), graphs
1070 return " ".join(content)
1073def printable_dim(dim: TensorShapeProto.Dimension) -> str:
1074 which = dim.WhichOneof("value")
1075 if which is None:
1076 return "?"
1077 return str(getattr(dim, which))
1080def printable_type(t: TypeProto) -> str:
1081 if t.WhichOneof("value") == "tensor_type":
1082 s: str = TensorProto.DataType.Name(t.tensor_type.elem_type) # type: ignore[arg-type]
1083 if t.tensor_type.HasField("shape"):
1084 if len(t.tensor_type.shape.dim):
1085 s += str(", " + "x".join(map(printable_dim, t.tensor_type.shape.dim)))
1086 else:
1087 s += ", scalar"
1088 return s
1089 if t.WhichOneof("value") is None:
1090 return ""
1091 return f"Unknown type {t.WhichOneof('value')}"
1094def printable_value_info(v: ValueInfoProto) -> str:
1095 s = f"%{v.name}"
1096 if v.type:
1097 s = f"{s}[{printable_type(v.type)}]"
1098 return s
1101def printable_tensor_proto(t: TensorProto) -> str:
1102 s = f"%{t.name}["
1103 s += TensorProto.DataType.Name(t.data_type) # type: ignore[arg-type]
1104 if t.dims is not None:
1105 if len(t.dims):
1106 s += str(", " + "x".join(map(str, t.dims)))
1107 else:
1108 s += ", scalar"
1109 s += "]"
1110 return s
1113def printable_node(
1114 node: NodeProto, prefix: str = "", subgraphs: bool = False
1115) -> str | tuple[str, list[GraphProto]]:
1116 content = []
1117 if len(node.output):
1118 content.append(", ".join([f"%{name}" for name in node.output]))
1119 content.append("=")
1120 # To deal with nested graphs
1121 graphs: list[GraphProto] = []
1122 printed_attrs = []
1123 for attr in node.attribute:
1124 if subgraphs:
1125 printed_attr_subgraphs = printable_attribute(attr, subgraphs)
1126 if not isinstance(printed_attr_subgraphs[1], list):
1127 raise TypeError(
1128 f"printed_attr_subgraphs[1] must be an instance of {list}."
1129 )
1130 graphs.extend(printed_attr_subgraphs[1])
1131 printed_attrs.append(printed_attr_subgraphs[0])
1132 else:
1133 printed = printable_attribute(attr)
1134 if not isinstance(printed, str):
1135 raise TypeError(f"printed must be an instance of {str}.")
1136 printed_attrs.append(printed)
1137 printed_attributes = ", ".join(sorted(printed_attrs))
1138 printed_inputs = ", ".join([f"%{name}" for name in node.input])
1139 if node.attribute:
1140 content.append(f"{node.op_type}[{printed_attributes}]({printed_inputs})")
1141 else:
1142 content.append(f"{node.op_type}({printed_inputs})")
1143 if subgraphs:
1144 return prefix + " ".join(content), graphs
1145 return prefix + " ".join(content)
1148@typing_extensions.deprecated(
1149 "Deprecated since 1.19. Consider using onnx.printer.to_text() instead."
1150)
1151def printable_graph(graph: GraphProto, prefix: str = "") -> str:
1152 """Display a GraphProto as a string.
1154 .. deprecated:: 1.19
1155 Consider using :func:`onnx.printer.to_text` instead.
1157 Args:
1158 graph (GraphProto): the graph to display
1159 prefix (string): prefix of every line
1161 Returns:
1162 string
1163 """
1164 content = []
1165 indent = prefix + " "
1166 # header
1167 header = ["graph", graph.name]
1168 initializers = {t.name for t in graph.initializer}
1169 if len(graph.input):
1170 header.append("(")
1171 in_strs = [] # required inputs
1172 in_with_init_strs: list = [] # optional inputs with initializer providing default value
1173 for inp in graph.input:
1174 if inp.name not in initializers:
1175 in_strs.append(printable_value_info(inp))
1176 else:
1177 in_with_init_strs.append(printable_value_info(inp))
1178 if in_strs:
1179 content.append(prefix + " ".join(header))
1180 header = []
1181 for line in in_strs:
1182 content.append(prefix + " " + line) # noqa: PERF401
1183 header.append(")")
1185 if in_with_init_strs:
1186 header.append("optional inputs with matching initializers (")
1187 content.append(prefix + " ".join(header))
1188 header = []
1189 for line in in_with_init_strs:
1190 content.append(prefix + " " + line) # noqa: PERF401
1191 header.append(")")
1193 # from IR 4 onwards an initializer is not required to have a matching graph input
1194 # so output the name, type and shape of those as well
1195 if len(in_with_init_strs) < len(initializers):
1196 graph_inputs = {i.name for i in graph.input}
1197 init_strs = [
1198 printable_tensor_proto(i)
1199 for i in graph.initializer
1200 if i.name not in graph_inputs
1201 ]
1202 header.append("initializers (")
1203 content.append(prefix + " ".join(header))
1204 header = []
1205 for line in init_strs:
1206 content.append(prefix + " " + line) # noqa: PERF401
1207 header.append(")")
1209 header.append("{")
1210 content.append(prefix + " ".join(header))
1211 graphs: list[GraphProto] = []
1212 # body
1213 for node in graph.node:
1214 contents_subgraphs = printable_node(node, indent, subgraphs=True)
1215 if not isinstance(contents_subgraphs[1], list):
1216 raise TypeError(f"contents_subgraphs[1] must be an instance of {list}.")
1217 content.append(contents_subgraphs[0])
1218 graphs.extend(contents_subgraphs[1])
1219 # tail
1220 tail = ["return"]
1221 if len(graph.output):
1222 tail.append(", ".join([f"%{out.name}" for out in graph.output]))
1223 content.append(indent + " ".join(tail))
1224 # closing bracket
1225 content.append(prefix + "}")
1226 for g in graphs:
1227 content.append("\n" + printable_graph(g)) # noqa: PERF401
1228 return "\n".join(content)
1231def strip_doc_string(proto: google.protobuf.message.Message) -> None:
1232 """Empties `doc_string` field on any nested protobuf messages"""
1233 if not isinstance(proto, google.protobuf.message.Message):
1234 raise TypeError(
1235 f"proto must be an instance of {google.protobuf.message.Message}."
1236 )
1237 for descriptor in proto.DESCRIPTOR.fields:
1238 if descriptor.name == "doc_string":
1239 proto.ClearField(descriptor.name)
1240 elif descriptor.type == descriptor.TYPE_MESSAGE:
1241 if descriptor.label == descriptor.LABEL_REPEATED:
1242 for x in getattr(proto, descriptor.name):
1243 strip_doc_string(x)
1244 elif proto.HasField(descriptor.name):
1245 strip_doc_string(getattr(proto, descriptor.name))
1248def make_training_info(
1249 algorithm: GraphProto,
1250 algorithm_bindings: AssignmentBindingType,
1251 initialization: GraphProto | None,
1252 initialization_bindings: AssignmentBindingType | None,
1253) -> TrainingInfoProto:
1254 training_info = TrainingInfoProto()
1255 training_info.algorithm.CopyFrom(algorithm)
1256 for k, v in algorithm_bindings:
1257 binding = training_info.update_binding.add()
1258 binding.key = k
1259 binding.value = v
1261 if initialization:
1262 training_info.initialization.CopyFrom(initialization)
1263 if initialization_bindings:
1264 for k, v in initialization_bindings:
1265 binding = training_info.initialization_binding.add()
1266 binding.key = k
1267 binding.value = v
1269 return training_info
1272# Following functions are used for mapping
1273def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
1274 """Convert a TensorProto's data_type to corresponding numpy dtype. It can be used while making tensor.
1276 Args:
1277 tensor_dtype: TensorProto's data_type
1279 Returns:
1280 numpy's data_type
1281 """
1282 return _mapping.TENSOR_TYPE_MAP[tensor_dtype].np_dtype
1285def tensor_dtype_to_storage_tensor_dtype(tensor_dtype: int) -> int:
1286 """Convert a TensorProto's data_type to corresponding data_type for storage.
1288 Args:
1289 tensor_dtype: TensorProto's data_type
1291 Returns:
1292 data_type for storage
1293 """
1294 return _mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
1297def tensor_dtype_to_string(tensor_dtype: int) -> str:
1298 """Get the name of given TensorProto's data_type.
1300 Args:
1301 tensor_dtype: TensorProto's data_type
1303 Returns:
1304 the name of data_type
1305 """
1306 return _mapping.TENSOR_TYPE_MAP[tensor_dtype].name
1309@functools.lru_cache(None)
1310def tensor_dtype_to_field(tensor_dtype: int) -> str:
1311 """Convert a TensorProto's data_type to corresponding field name for storage. It can be used while making tensors.
1313 Args:
1314 tensor_dtype: TensorProto's data_type
1316 Returns:
1317 field name
1318 """
1319 storage_tensor_type_to_field = {
1320 int(TensorProto.FLOAT): "float_data",
1321 int(TensorProto.INT32): "int32_data",
1322 int(TensorProto.INT64): "int64_data",
1323 int(TensorProto.DOUBLE): "double_data",
1324 int(TensorProto.UINT32): "uint64_data",
1325 int(TensorProto.UINT64): "uint64_data",
1326 int(TensorProto.STRING): "string_data",
1327 }
1328 return storage_tensor_type_to_field[
1329 _mapping.TENSOR_TYPE_MAP[tensor_dtype].storage_dtype
1330 ]
1333@functools.lru_cache(None)
1334def np_dtype_to_tensor_dtype(np_dtype: np.dtype) -> TensorProto.DataType:
1335 """Convert a numpy's dtype to corresponding tensor type. It can be used while converting numpy arrays to tensors.
1337 Args:
1338 np_dtype: numpy's data_type
1340 Returns:
1341 TensorsProto's data_type
1342 """
1343 _np_dtype_to_tensor_dtype = {
1344 v.np_dtype: k for k, v in _mapping.TENSOR_TYPE_MAP.items()
1345 }
1346 if np_dtype in _np_dtype_to_tensor_dtype:
1347 return typing.cast("TensorProto.DataType", _np_dtype_to_tensor_dtype[np_dtype])
1348 if np.issubdtype(np_dtype, np.str_):
1349 return TensorProto.STRING # type: ignore[return-value]
1351 raise ValueError(
1352 f"Unable to convert type {np_dtype!r} into TensorProto element type."
1353 )
1356def get_all_tensor_dtypes() -> KeysView[int]:
1357 """Get all tensor types from TensorProto.
1359 Returns:
1360 all tensor types from TensorProto
1361 """
1362 return _mapping.TENSOR_TYPE_MAP.keys()
1365_ATTRIBUTE_TYPE_TO_STR: dict[int, str] = {
1366 k: v for v, k in AttributeProto.AttributeType.items()
1367}
1370def _attr_type_to_str(attr_type: int) -> str:
1371 """Convert AttributeProto type to string.
1373 Args:
1374 attr_type: AttributeProto type.
1376 Returns:
1377 String representing the supplied attr_type.
1378 """
1379 if attr_type in AttributeProto.AttributeType.values():
1380 return _ATTRIBUTE_TYPE_TO_STR[attr_type]
1381 return AttributeProto.AttributeType.keys()[0]