onnx-ir 0.1.6__tar.gz → 0.1.7__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of onnx-ir might be problematic. Click here for more details.

Files changed (53) hide show
  1. {onnx_ir-0.1.6/src/onnx_ir.egg-info → onnx_ir-0.1.7}/PKG-INFO +2 -3
  2. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/pyproject.toml +3 -3
  3. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/__init__.py +1 -1
  4. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_convenience/__init__.py +49 -32
  5. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_core.py +60 -16
  6. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/__init__.py +2 -0
  7. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/constant_manipulation.py +15 -7
  8. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/identity_elimination.py +1 -0
  9. onnx_ir-0.1.7/src/onnx_ir/passes/common/initializer_deduplication.py +167 -0
  10. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/naming.py +1 -1
  11. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/serde.py +94 -34
  12. {onnx_ir-0.1.6 → onnx_ir-0.1.7/src/onnx_ir.egg-info}/PKG-INFO +2 -3
  13. onnx_ir-0.1.6/src/onnx_ir/passes/common/initializer_deduplication.py +0 -56
  14. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/LICENSE +0 -0
  15. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/MANIFEST.in +0 -0
  16. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/README.md +0 -0
  17. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/setup.cfg +0 -0
  18. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_convenience/_constructors.py +0 -0
  19. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_display.py +0 -0
  20. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_enums.py +0 -0
  21. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_graph_comparison.py +0 -0
  22. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_graph_containers.py +0 -0
  23. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_io.py +0 -0
  24. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_linked_list.py +0 -0
  25. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_metadata.py +0 -0
  26. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_name_authority.py +0 -0
  27. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_polyfill.py +0 -0
  28. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_protocols.py +0 -0
  29. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_tape.py +0 -0
  30. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_thirdparty/asciichartpy.py +0 -0
  31. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_type_casting.py +0 -0
  32. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/_version_utils.py +0 -0
  33. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/convenience.py +0 -0
  34. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/external_data.py +0 -0
  35. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/__init__.py +0 -0
  36. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/_pass_infra.py +0 -0
  37. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/_c_api_utils.py +0 -0
  38. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/clear_metadata_and_docstring.py +0 -0
  39. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/common_subexpression_elimination.py +0 -0
  40. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/inliner.py +0 -0
  41. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/onnx_checker.py +0 -0
  42. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/shape_inference.py +0 -0
  43. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/topological_sort.py +0 -0
  44. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/passes/common/unused_removal.py +0 -0
  45. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/py.typed +0 -0
  46. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/tape.py +0 -0
  47. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/tensor_adapters.py +0 -0
  48. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/testing.py +0 -0
  49. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir/traversal.py +0 -0
  50. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir.egg-info/SOURCES.txt +0 -0
  51. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir.egg-info/dependency_links.txt +0 -0
  52. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir.egg-info/requires.txt +0 -0
  53. {onnx_ir-0.1.6 → onnx_ir-0.1.7}/src/onnx_ir.egg-info/top_level.txt +0 -0
@@ -1,9 +1,9 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: onnx-ir
3
- Version: 0.1.6
3
+ Version: 0.1.7
4
4
  Summary: Efficient in-memory representation for ONNX
5
5
  Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
6
- License: Apache License v2.0
6
+ License-Expression: Apache-2.0
7
7
  Project-URL: Homepage, https://onnx.ai/ir-py
8
8
  Project-URL: Issues, https://github.com/onnx/ir-py/issues
9
9
  Project-URL: Repository, https://github.com/onnx/ir-py
@@ -13,7 +13,6 @@ Classifier: Programming Language :: Python :: 3.10
13
13
  Classifier: Programming Language :: Python :: 3.11
14
14
  Classifier: Programming Language :: Python :: 3.12
15
15
  Classifier: Programming Language :: Python :: 3.13
16
- Classifier: License :: OSI Approved :: Apache Software License
17
16
  Requires-Python: >=3.9
18
17
  Description-Content-Type: text/markdown
19
18
  License-File: LICENSE
@@ -1,5 +1,5 @@
1
1
  [build-system]
2
- requires = ["setuptools>=70"]
2
+ requires = ["setuptools>=77"]
3
3
  build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
@@ -11,7 +11,8 @@ authors = [
11
11
  ]
12
12
  readme = "README.md"
13
13
  requires-python = ">=3.9"
14
- license = {text = "Apache License v2.0"}
14
+ license = "Apache-2.0"
15
+ license-files = ["LICEN[CS]E*"]
15
16
  classifiers = [
16
17
  "Development Status :: 4 - Beta",
17
18
  "Programming Language :: Python :: 3.9",
@@ -19,7 +20,6 @@ classifiers = [
19
20
  "Programming Language :: Python :: 3.11",
20
21
  "Programming Language :: Python :: 3.12",
21
22
  "Programming Language :: Python :: 3.13",
22
- "License :: OSI Approved :: Apache Software License",
23
23
  ]
24
24
  dependencies = ["numpy", "onnx>=1.16", "typing_extensions>=4.10", "ml_dtypes"]
25
25
 
@@ -167,4 +167,4 @@ def __set_module() -> None:
167
167
 
168
168
 
169
169
  __set_module()
170
- __version__ = "0.1.6"
170
+ __version__ = "0.1.7"
@@ -58,44 +58,52 @@ def _infer_attribute_type(attr: SupportedAttrTypes) -> _enums.AttributeType:
58
58
  return _enums.AttributeType.STRING
59
59
  if isinstance(attr, _core.Attr):
60
60
  return attr.type
61
- if isinstance(attr, Sequence) and all(isinstance(x, int) for x in attr):
62
- return _enums.AttributeType.INTS
63
- if isinstance(attr, Sequence) and all(isinstance(x, float) for x in attr):
64
- return _enums.AttributeType.FLOATS
65
- if isinstance(attr, Sequence) and all(isinstance(x, str) for x in attr):
66
- return _enums.AttributeType.STRINGS
61
+ if isinstance(attr, (_core.Graph, onnx.GraphProto, _protocols.GraphProtocol)):
62
+ return _enums.AttributeType.GRAPH
67
63
  if isinstance(attr, (_core.TensorBase, onnx.TensorProto, _protocols.TensorProtocol)):
68
64
  # Be sure to check TensorProtocol last because isinstance checking on Protocols can be slower
69
65
  return _enums.AttributeType.TENSOR
70
- if isinstance(attr, Sequence) and all(
71
- isinstance(x, (_core.TensorBase, onnx.TensorProto, _protocols.TensorProtocol))
72
- for x in attr
73
- ):
74
- return _enums.AttributeType.TENSORS
75
- if isinstance(attr, (_core.Graph, onnx.GraphProto, _protocols.GraphProtocol)):
76
- return _enums.AttributeType.GRAPH
77
- if isinstance(attr, Sequence) and all(
78
- isinstance(x, (_core.Graph, onnx.GraphProto, _protocols.GraphProtocol)) for x in attr
79
- ):
80
- return _enums.AttributeType.GRAPHS
81
66
  if isinstance(
82
67
  attr,
83
68
  (_core.TensorType, _core.SequenceType, _core.OptionalType, _protocols.TypeProtocol),
84
69
  ):
85
70
  return _enums.AttributeType.TYPE_PROTO
86
- if isinstance(attr, Sequence) and all(
87
- isinstance(
88
- x,
89
- (
90
- _core.TensorType,
91
- _core.SequenceType,
92
- _core.OptionalType,
93
- _protocols.TypeProtocol,
94
- ),
95
- )
96
- for x in attr
97
- ):
98
- return _enums.AttributeType.TYPE_PROTOS
71
+ if isinstance(attr, Sequence):
72
+ if not attr:
73
+ logger.warning(
74
+ "Attribute type is ambiguous because it is an empty sequence. "
75
+ "Please create an Attr with an explicit type. Defaulted to INTS"
76
+ )
77
+ return _enums.AttributeType.INTS
78
+ if all(isinstance(x, int) for x in attr):
79
+ return _enums.AttributeType.INTS
80
+ if all(isinstance(x, float) for x in attr):
81
+ return _enums.AttributeType.FLOATS
82
+ if all(isinstance(x, str) for x in attr):
83
+ return _enums.AttributeType.STRINGS
84
+ if all(
85
+ isinstance(x, (_core.TensorBase, onnx.TensorProto, _protocols.TensorProtocol))
86
+ for x in attr
87
+ ):
88
+ return _enums.AttributeType.TENSORS
89
+ if all(
90
+ isinstance(x, (_core.Graph, onnx.GraphProto, _protocols.GraphProtocol))
91
+ for x in attr
92
+ ):
93
+ return _enums.AttributeType.GRAPHS
94
+ if all(
95
+ isinstance(
96
+ x,
97
+ (
98
+ _core.TensorType,
99
+ _core.SequenceType,
100
+ _core.OptionalType,
101
+ _protocols.TypeProtocol,
102
+ ),
103
+ )
104
+ for x in attr
105
+ ):
106
+ return _enums.AttributeType.TYPE_PROTOS
99
107
  raise TypeError(f"Unsupported attribute type: '{type(attr)}'")
100
108
 
101
109
 
@@ -218,7 +226,7 @@ def convert_attributes(
218
226
  ... "type_protos": [ir.TensorType(ir.DataType.FLOAT), ir.TensorType(ir.DataType.FLOAT)],
219
227
  ... }
220
228
  >>> convert_attributes(attrs)
221
- [Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', INTS, Graph(
229
+ [Attr('int', INT, 1), Attr('float', FLOAT, 1.0), Attr('str', STRING, 'hello'), Attr('ints', INTS, [1, 2, 3]), Attr('floats', FLOATS, [1.0, 2.0, 3.0]), Attr('strings', STRINGS, ['hello', 'world']), Attr('tensor', TENSOR, Tensor<DOUBLE,[3]>(array([1., 2., 3.]), name=None)), Attr('tensor_proto', TENSOR, TensorProtoTensor<FLOAT,[3]>(array([1., 2., 3.], dtype=float32), name='proto')), Attr('graph', GRAPH, Graph(
222
230
  name='graph0',
223
231
  inputs=(
224
232
  <BLANKLINE>
@@ -247,11 +255,20 @@ def convert_attributes(
247
255
  len()=0
248
256
  )]), Attr('type_proto', TYPE_PROTO, Tensor(FLOAT)), Attr('type_protos', TYPE_PROTOS, [Tensor(FLOAT), Tensor(FLOAT)])]
249
257
 
258
+ .. important::
259
+ An empty sequence should be created with an explicit type by initializing
260
+ an Attr object with an attribute type to avoid type ambiguity. For example::
261
+
262
+ ir.Attr("empty", [], type=ir.AttributeType.INTS)
263
+
250
264
  Args:
251
265
  attrs: A dictionary of {<attribute name>: <python objects>} to convert.
252
266
 
253
267
  Returns:
254
- A list of _core.Attr objects.
268
+ A list of :class:`_core.Attr` objects.
269
+
270
+ Raises:
271
+ TypeError: If an attribute type is not supported.
255
272
  """
256
273
  attributes: list[_core.Attr] = []
257
274
  for name, attr in attrs.items():
@@ -2564,14 +2564,23 @@ class Graph(_protocols.GraphProtocol, Sequence[Node], _display.PrettyPrintable):
2564
2564
 
2565
2565
  .. versionadded:: 0.1.2
2566
2566
  """
2567
- seen_graphs: set[Graph] = set()
2568
- for node in onnx_ir.traversal.RecursiveGraphIterator(self):
2569
- graph = node.graph
2567
+ # Use a dict to preserve order
2568
+ seen_graphs: dict[Graph, None] = {}
2569
+
2570
+ # Need to use the enter_graph callback so that empty subgraphs are collected
2571
+ def enter_subgraph(graph) -> None:
2570
2572
  if graph is self:
2571
- continue
2572
- if graph is not None and graph not in seen_graphs:
2573
- seen_graphs.add(graph)
2574
- yield graph
2573
+ return
2574
+ if not isinstance(graph, Graph):
2575
+ raise TypeError(
2576
+ f"Expected a Graph, got {type(graph)}. The model may be invalid"
2577
+ )
2578
+ if graph not in seen_graphs:
2579
+ seen_graphs[graph] = None
2580
+
2581
+ for _ in onnx_ir.traversal.RecursiveGraphIterator(self, enter_graph=enter_subgraph):
2582
+ pass
2583
+ yield from seen_graphs.keys()
2575
2584
 
2576
2585
  # Mutation methods
2577
2586
  def append(self, node: Node, /) -> None:
@@ -3180,6 +3189,21 @@ class Function(_protocols.FunctionProtocol, Sequence[Node], _display.PrettyPrint
3180
3189
  def attributes(self) -> _graph_containers.Attributes:
3181
3190
  return self._attributes
3182
3191
 
3192
+ @property
3193
+ def graph(self) -> Graph:
3194
+ """The underlying Graph object that contains the nodes of this function.
3195
+
3196
+ Only use this graph for identity comparison::
3197
+
3198
+ if value.graph is function.graph:
3199
+ # Do something with the value that belongs to this function
3200
+
3201
+ Otherwise use the Function object directly to access the nodes and other properties.
3202
+
3203
+ .. versionadded:: 0.1.7
3204
+ """
3205
+ return self._graph
3206
+
3183
3207
  @typing.overload
3184
3208
  def __getitem__(self, index: int) -> Node: ...
3185
3209
  @typing.overload
@@ -3240,14 +3264,22 @@ class Function(_protocols.FunctionProtocol, Sequence[Node], _display.PrettyPrint
3240
3264
 
3241
3265
  .. versionadded:: 0.1.2
3242
3266
  """
3243
- seen_graphs: set[Graph] = set()
3244
- for node in onnx_ir.traversal.RecursiveGraphIterator(self):
3245
- graph = node.graph
3246
- if graph is self._graph:
3247
- continue
3248
- if graph is not None and graph not in seen_graphs:
3249
- seen_graphs.add(graph)
3250
- yield graph
3267
+ seen_graphs: dict[Graph, None] = {}
3268
+
3269
+ # Need to use the enter_graph callback so that empty subgraphs are collected
3270
+ def enter_subgraph(graph) -> None:
3271
+ if graph is self:
3272
+ return
3273
+ if not isinstance(graph, Graph):
3274
+ raise TypeError(
3275
+ f"Expected a Graph, got {type(graph)}. The model may be invalid"
3276
+ )
3277
+ if graph not in seen_graphs:
3278
+ seen_graphs[graph] = None
3279
+
3280
+ for _ in onnx_ir.traversal.RecursiveGraphIterator(self, enter_graph=enter_subgraph):
3281
+ pass
3282
+ yield from seen_graphs.keys()
3251
3283
 
3252
3284
  # Mutation methods
3253
3285
  def append(self, node: Node, /) -> None:
@@ -3349,7 +3381,7 @@ class Attr(
3349
3381
  ):
3350
3382
  """Base class for ONNX attributes or references."""
3351
3383
 
3352
- __slots__ = ("_name", "_ref_attr_name", "_type", "_value", "doc_string")
3384
+ __slots__ = ("_metadata", "_name", "_ref_attr_name", "_type", "_value", "doc_string")
3353
3385
 
3354
3386
  def __init__(
3355
3387
  self,
@@ -3365,6 +3397,7 @@ class Attr(
3365
3397
  self._value = value
3366
3398
  self._ref_attr_name = ref_attr_name
3367
3399
  self.doc_string = doc_string
3400
+ self._metadata: _metadata.MetadataStore | None = None
3368
3401
 
3369
3402
  @property
3370
3403
  def name(self) -> str:
@@ -3386,6 +3419,17 @@ class Attr(
3386
3419
  def ref_attr_name(self) -> str | None:
3387
3420
  return self._ref_attr_name
3388
3421
 
3422
+ @property
3423
+ def meta(self) -> _metadata.MetadataStore:
3424
+ """The metadata store for intermediate analysis.
3425
+
3426
+ Write to the :attr:`metadata_props` if you would like the metadata to be serialized
3427
+ to the ONNX proto.
3428
+ """
3429
+ if self._metadata is None:
3430
+ self._metadata = _metadata.MetadataStore()
3431
+ return self._metadata
3432
+
3389
3433
  def is_ref(self) -> bool:
3390
3434
  """Check if this attribute is a reference attribute."""
3391
3435
  return self.ref_attr_name is not None
@@ -6,6 +6,7 @@ __all__ = [
6
6
  "CheckerPass",
7
7
  "ClearMetadataAndDocStringPass",
8
8
  "CommonSubexpressionEliminationPass",
9
+ "DeduplicateHashedInitializersPass",
9
10
  "DeduplicateInitializersPass",
10
11
  "IdentityEliminationPass",
11
12
  "InlinePass",
@@ -36,6 +37,7 @@ from onnx_ir.passes.common.identity_elimination import (
36
37
  IdentityEliminationPass,
37
38
  )
38
39
  from onnx_ir.passes.common.initializer_deduplication import (
40
+ DeduplicateHashedInitializersPass,
39
41
  DeduplicateInitializersPass,
40
42
  )
41
43
  from onnx_ir.passes.common.inliner import InlinePass
@@ -148,6 +148,7 @@ class LiftSubgraphInitializersToMainGraphPass(ir.passes.InPlacePass):
148
148
  if graph is model.graph:
149
149
  continue
150
150
  for name in tuple(graph.initializers):
151
+ assert name is not None
151
152
  initializer = graph.initializers[name]
152
153
  if initializer.is_graph_input():
153
154
  # Skip the ones that are also graph inputs
@@ -156,17 +157,24 @@ class LiftSubgraphInitializersToMainGraphPass(ir.passes.InPlacePass):
156
157
  initializer.name,
157
158
  )
158
159
  continue
160
+ if initializer.is_graph_output():
161
+ logger.debug(
162
+ "Initializer '%s' is used as output, so it can't be lifted",
163
+ initializer.name,
164
+ )
165
+ continue
159
166
  # Remove the initializer from the subgraph
160
167
  graph.initializers.pop(name)
161
168
  # To avoid name conflicts, we need to rename the initializer
162
169
  # to a unique name in the main graph
163
- if name in registered_initializer_names:
164
- name_count = registered_initializer_names[name]
165
- initializer.name = f"{name}_{name_count}"
166
- registered_initializer_names[name] = name_count + 1
167
- else:
168
- assert initializer.name is not None
169
- registered_initializer_names[initializer.name] = 1
170
+ new_name = name
171
+ while new_name in model.graph.initializers:
172
+ if name in registered_initializer_names:
173
+ registered_initializer_names[name] += 1
174
+ else:
175
+ registered_initializer_names[name] = 1
176
+ new_name = f"{name}_{registered_initializer_names[name]}"
177
+ initializer.name = new_name
170
178
  model.graph.register_initializer(initializer)
171
179
  count += 1
172
180
  logger.debug(
@@ -19,6 +19,7 @@ class IdentityEliminationPass(ir.passes.InPlacePass):
19
19
  """Pass for eliminating redundant Identity nodes.
20
20
 
21
21
  This pass removes Identity nodes according to the following rules:
22
+
22
23
  1. For any node of the form `y = Identity(x)`, where `y` is not an output
23
24
  of any graph, replace all uses of `y` with a use of `x`, and remove the node.
24
25
  2. If `y` is an output of a graph, and `x` is not an input of any graph,
@@ -0,0 +1,167 @@
1
+ # Copyright (c) ONNX Project Contributors
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ """Pass for removing duplicated initializer tensors from a graph."""
4
+
5
+ from __future__ import annotations
6
+
7
+ __all__ = ["DeduplicateInitializersPass", "DeduplicateHashedInitializersPass"]
8
+
9
+
10
+ import hashlib
11
+ import logging
12
+
13
+ import onnx_ir as ir
14
+
15
+ logger = logging.getLogger(__name__)
16
+
17
+
18
+ def _should_skip_initializer(initializer: ir.Value, size_limit: int) -> bool:
19
+ """Check if the initializer should be skipped for deduplication."""
20
+ if initializer.is_graph_input() or initializer.is_graph_output():
21
+ # Skip graph inputs and outputs
22
+ logger.warning(
23
+ "Skipped deduplication of initializer '%s' as it is a graph input or output",
24
+ initializer.name,
25
+ )
26
+ return True
27
+
28
+ const_val = initializer.const_value
29
+ if const_val is None:
30
+ # Skip if initializer has no constant value
31
+ logger.warning(
32
+ "Skipped deduplication of initializer '%s' as it has no constant value. The model may contain invalid initializers",
33
+ initializer.name,
34
+ )
35
+ return True
36
+
37
+ if const_val.size > size_limit:
38
+ # Skip if the initializer is larger than the size limit
39
+ logger.debug(
40
+ "Skipped initializer '%s' as it exceeds the size limit of %d elements",
41
+ initializer.name,
42
+ size_limit,
43
+ )
44
+ return True
45
+
46
+ if const_val.dtype == ir.DataType.STRING:
47
+ # Skip string initializers as they don't have a bytes representation
48
+ logger.warning(
49
+ "Skipped deduplication of string initializer '%s' (unsupported yet)",
50
+ initializer.name,
51
+ )
52
+ return True
53
+ return False
54
+
55
+
56
+ class DeduplicateInitializersPass(ir.passes.InPlacePass):
57
+ """Remove duplicated initializer tensors from the main graph and all subgraphs.
58
+
59
+ This pass detects initializers with identical shape, dtype, and content,
60
+ and replaces all duplicate references with a canonical one.
61
+
62
+ Initializers are deduplicated within each graph. To deduplicate initializers
63
+ in the model globally (across graphs), use :class:`~onnx_ir.passes.common.LiftSubgraphInitializersToMainGraphPass`
64
+ to lift the initializers to the main graph first before running pass.
65
+
66
+ .. versionadded:: 0.1.3
67
+ .. versionchanged:: 0.1.7
68
+ This pass now deduplicates initializers in subgraphs as well.
69
+ """
70
+
71
+ def __init__(self, size_limit: int = 1024):
72
+ super().__init__()
73
+ self.size_limit = size_limit
74
+
75
+ def call(self, model: ir.Model) -> ir.passes.PassResult:
76
+ modified = False
77
+
78
+ for graph in model.graphs():
79
+ initializers: dict[tuple[ir.DataType, tuple[int, ...], bytes], ir.Value] = {}
80
+ for initializer in tuple(graph.initializers.values()):
81
+ if _should_skip_initializer(initializer, self.size_limit):
82
+ continue
83
+
84
+ const_val = initializer.const_value
85
+ assert const_val is not None
86
+
87
+ key = (const_val.dtype, tuple(const_val.shape), const_val.tobytes())
88
+ if key in initializers:
89
+ modified = True
90
+ initializer_to_keep = initializers[key] # type: ignore[index]
91
+ ir.convenience.replace_all_uses_with(initializer, initializer_to_keep)
92
+ assert initializer.name is not None
93
+ graph.initializers.pop(initializer.name)
94
+ logger.info(
95
+ "Replaced initializer '%s' with existing initializer '%s'",
96
+ initializer.name,
97
+ initializer_to_keep.name,
98
+ )
99
+ else:
100
+ initializers[key] = initializer # type: ignore[index]
101
+
102
+ return ir.passes.PassResult(model=model, modified=modified)
103
+
104
+
105
+ class DeduplicateHashedInitializersPass(ir.passes.InPlacePass):
106
+ """Remove duplicated initializer tensors (using a hashed method) from the graph.
107
+
108
+ This pass detects initializers with identical shape, dtype, and hashed content,
109
+ and replaces all duplicate references with a canonical one.
110
+
111
+ This pass should have a lower peak memory usage than :class:`DeduplicateInitializersPass`
112
+ as it does not store the full tensor data in memory, but instead uses a hash of the tensor data.
113
+
114
+ .. versionadded:: 0.1.7
115
+ """
116
+
117
+ def __init__(self, size_limit: int = 4 * 1024 * 1024 * 1024):
118
+ super().__init__()
119
+ # 4 GB default size limit for deduplication
120
+ self.size_limit = size_limit
121
+
122
+ def call(self, model: ir.Model) -> ir.passes.PassResult:
123
+ modified = False
124
+
125
+ for graph in model.graphs():
126
+ initializers: dict[tuple[ir.DataType, tuple[int, ...], str], ir.Value] = {}
127
+
128
+ for initializer in tuple(graph.initializers.values()):
129
+ if _should_skip_initializer(initializer, self.size_limit):
130
+ continue
131
+
132
+ const_val = initializer.const_value
133
+ assert const_val is not None
134
+
135
+ # Hash tensor data to avoid storing large amounts of data in memory
136
+ hashed = hashlib.sha512()
137
+ tensor_data = const_val.numpy()
138
+ hashed.update(tensor_data)
139
+ tensor_digest = hashed.hexdigest()
140
+
141
+ tensor_dims = tuple(const_val.shape.numpy())
142
+
143
+ key = (const_val.dtype, tensor_dims, tensor_digest)
144
+
145
+ if key in initializers:
146
+ if initializers[key].const_value.tobytes() != const_val.tobytes():
147
+ logger.warning(
148
+ "Initializer deduplication failed: "
149
+ "hashes match but values differ with values %s and %s",
150
+ initializers[key],
151
+ initializer,
152
+ )
153
+ continue
154
+ modified = True
155
+ initializer_to_keep = initializers[key] # type: ignore[index]
156
+ ir.convenience.replace_all_uses_with(initializer, initializer_to_keep)
157
+ assert initializer.name is not None
158
+ graph.initializers.pop(initializer.name)
159
+ logger.info(
160
+ "Replaced initializer '%s' with existing initializer '%s'",
161
+ initializer.name,
162
+ initializer_to_keep.name,
163
+ )
164
+ else:
165
+ initializers[key] = initializer # type: ignore[index]
166
+
167
+ return ir.passes.PassResult(model=model, modified=modified)
@@ -64,7 +64,7 @@ class NameFixPass(ir.passes.InPlacePass):
64
64
  def custom_value_name(value: ir.Value) -> str:
65
65
  return f"custom_value_{value.type}"
66
66
 
67
- name_fix_pass = NameFixPass(nameGenerator=CustomNameGenerator())
67
+ name_fix_pass = NameFixPass(name_generator=CustomNameGenerator())
68
68
 
69
69
  .. versionadded:: 0.1.6
70
70
  """
@@ -682,8 +682,8 @@ def deserialize_graph(proto: onnx.GraphProto) -> _core.Graph:
682
682
  Returns:
683
683
  IR Graph.
684
684
 
685
- .. versionadded:: 0.3
686
- Support for *quantization_annotation* is added.
685
+ .. versionadded:: 0.1.3
686
+ Support for `quantization_annotation` is added.
687
687
  """
688
688
  return _deserialize_graph(proto, [])
689
689
 
@@ -760,6 +760,18 @@ def _deserialize_graph(
760
760
  # Build the value info dictionary to allow for quick lookup for this graph scope
761
761
  value_info = {info.name: info for info in proto.value_info}
762
762
 
763
+ # Declare values for all node outputs from this graph scope. This is necessary
764
+ # to handle the case where a node in a subgraph uses a value that is declared out
765
+ # of order in the outer graph. Declaring the values first allows us to find the
766
+ # values later when deserializing the nodes in subgraphs.
767
+ for node in proto.node:
768
+ _declare_node_outputs(
769
+ node,
770
+ values,
771
+ value_info=value_info,
772
+ quantization_annotations=quantization_annotations,
773
+ )
774
+
763
775
  # Deserialize nodes with all known values
764
776
  nodes = [
765
777
  _deserialize_node(node, scoped_values, value_info, quantization_annotations)
@@ -798,6 +810,55 @@ def _deserialize_graph(
798
810
  )
799
811
 
800
812
 
813
+ def _declare_node_outputs(
814
+ proto: onnx.NodeProto,
815
+ current_value_scope: dict[str, _core.Value],
816
+ value_info: dict[str, onnx.ValueInfoProto],
817
+ quantization_annotations: dict[str, onnx.TensorAnnotation],
818
+ ) -> None:
819
+ """Declare outputs for a node in the current graph scope.
820
+
821
+ This is necessary to handle the case where a node in a subgraph uses a value that is declared
822
+ out of order in the outer graph. Declaring the values first allows us to find the values later
823
+ when deserializing the nodes in subgraphs.
824
+
825
+ Args:
826
+ proto: The ONNX NodeProto to declare outputs for.
827
+ current_value_scope: The current scope of values, mapping value names to their corresponding Value objects.
828
+ value_info: A dictionary mapping value names to their corresponding ValueInfoProto.
829
+ quantization_annotations: A dictionary mapping tensor names to their corresponding TensorAnnotation.
830
+
831
+ Raises:
832
+ ValueError: If an output name is redeclared in the current graph scope.
833
+ """
834
+ for output_name in proto.output:
835
+ if output_name == "":
836
+ continue
837
+ if output_name in current_value_scope:
838
+ raise ValueError(
839
+ f"Output '{output_name}' is redeclared in the current graph scope. "
840
+ f"Original declaration {current_value_scope[output_name]}. "
841
+ f"New declaration: by operator '{proto.op_type}' of node '{proto.name}'. "
842
+ "The model is invalid"
843
+ )
844
+
845
+ # Create the value and add it to the current scope.
846
+ value = _core.Value(name=output_name)
847
+ current_value_scope[output_name] = value
848
+ # Fill in shape/type information if they exist
849
+ if output_name in value_info:
850
+ deserialize_value_info_proto(value_info[output_name], value)
851
+ else:
852
+ logger.debug(
853
+ "ValueInfoProto not found for output '%s' in node '%s' of type '%s'",
854
+ output_name,
855
+ proto.name,
856
+ proto.op_type,
857
+ )
858
+ if output_name in quantization_annotations:
859
+ _deserialize_quantization_annotation(quantization_annotations[output_name], value)
860
+
861
+
801
862
  @_capture_errors(lambda proto: proto.name)
802
863
  def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
803
864
  """Deserialize an ONNX FunctionProto into an IR Function.
@@ -812,7 +873,14 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
812
873
  values: dict[str, _core.Value] = {v.name: v for v in inputs} # type: ignore[misc]
813
874
  value_info = {info.name: info for info in getattr(proto, "value_info", [])}
814
875
 
815
- # TODO(justinchuby): Handle unsorted nodes
876
+ for node in proto.node:
877
+ _declare_node_outputs(
878
+ node,
879
+ values,
880
+ value_info=value_info,
881
+ quantization_annotations={},
882
+ )
883
+
816
884
  nodes = [
817
885
  _deserialize_node(node, [values], value_info=value_info, quantization_annotations={})
818
886
  for node in proto.node
@@ -1137,8 +1205,15 @@ def deserialize_node(proto: onnx.NodeProto) -> _core.Node:
1137
1205
  Returns:
1138
1206
  An IR Node object representing the ONNX node.
1139
1207
  """
1208
+ value_scope: dict[str, _core.Value] = {}
1209
+ _declare_node_outputs(
1210
+ proto,
1211
+ value_scope,
1212
+ value_info={},
1213
+ quantization_annotations={},
1214
+ )
1140
1215
  return _deserialize_node(
1141
- proto, scoped_values=[{}], value_info={}, quantization_annotations={}
1216
+ proto, scoped_values=[value_scope], value_info={}, quantization_annotations={}
1142
1217
  )
1143
1218
 
1144
1219
 
@@ -1161,18 +1236,18 @@ def _deserialize_node(
1161
1236
  for values in reversed(scoped_values):
1162
1237
  if input_name not in values:
1163
1238
  continue
1239
+
1164
1240
  node_inputs.append(values[input_name])
1165
1241
  found = True
1166
1242
  del values # Remove the reference so it is not used by mistake
1167
1243
  break
1168
1244
  if not found:
1169
- # If the input is not found, we know the graph may be unsorted and
1170
- # the input may be a supposed-to-be initializer or an output of a node that comes later.
1171
- # Here we create the value with the name and add it to the current scope.
1172
- # Nodes need to check the value pool for potentially initialized outputs
1245
+ # If the input is not found, we know the graph is invalid because the value
1246
+ # is not declared. We will still create a new input for the node so that
1247
+ # it can be fixed later.
1173
1248
  logger.warning(
1174
- "Input '%s' of node '%s(%s::%s:%s)' not found in any scope. "
1175
- "The graph may be unsorted. Creating a new input (current depth: %s) .",
1249
+ "Input '%s' of node '%s' (%s::%s:%s) cannot be found in any scope. "
1250
+ "The model is invalid but we will still create a new input for the node (current depth: %s)",
1176
1251
  input_name,
1177
1252
  proto.name,
1178
1253
  proto.domain,
@@ -1208,35 +1283,22 @@ def _deserialize_node(
1208
1283
  node_outputs.append(_core.Value(name=""))
1209
1284
  continue
1210
1285
 
1211
- # 1. When the graph is unsorted, we may be able to find the output already created
1286
+ # The outputs should already be declared in the current scope by _declare_node_outputs.
1287
+ #
1288
+ # When the graph is unsorted, we may be able to find the output already created
1212
1289
  # as an input to some other nodes in the current scope.
1213
1290
  # Note that a value is always owned by the producing node. Even though a value
1214
1291
  # can be created when parsing inputs of other nodes, the new node created here
1215
1292
  # that produces the value will assume ownership. It is then impossible to transfer
1216
1293
  # the ownership to any other node.
1217
-
1294
+ #
1218
1295
  # The output can only be found in the current scope. It is impossible for
1219
1296
  # a node to produce an output that is not in its own scope.
1220
1297
  current_scope = scoped_values[-1]
1221
- if output_name in current_scope:
1222
- value = current_scope[output_name]
1223
- else:
1224
- # 2. Common scenario: the graph is sorted and this is the first time we see the output.
1225
- # Create the value and add it to the current scope.
1226
- value = _core.Value(name=output_name)
1227
- current_scope[output_name] = value
1228
- # Fill in shape/type information if they exist
1229
- if output_name in value_info:
1230
- deserialize_value_info_proto(value_info[output_name], value)
1231
- else:
1232
- logger.debug(
1233
- "ValueInfoProto not found for output '%s' in node '%s' of type '%s'",
1234
- output_name,
1235
- proto.name,
1236
- proto.op_type,
1237
- )
1238
- if output_name in quantization_annotations:
1239
- _deserialize_quantization_annotation(quantization_annotations[output_name], value)
1298
+ assert output_name in current_scope, (
1299
+ f"Output '{output_name}' not found in the current scope. This is unexpected"
1300
+ )
1301
+ value = current_scope[output_name]
1240
1302
  node_outputs.append(value)
1241
1303
  return _core.Node(
1242
1304
  proto.domain,
@@ -1469,8 +1531,6 @@ def serialize_graph_into(
1469
1531
  serialize_value_into(graph_proto.input.add(), input_)
1470
1532
  if input_.name not in from_.initializers:
1471
1533
  # Annotations for initializers will be added below to avoid double adding
1472
- # TODO(justinchuby): We should add a method is_initializer() on Value when
1473
- # the initializer list is tracked
1474
1534
  _maybe_add_quantization_annotation(graph_proto, input_)
1475
1535
  input_names = {input_.name for input_ in from_.inputs}
1476
1536
  # TODO(justinchuby): Support sparse_initializer
@@ -1818,7 +1878,7 @@ def serialize_value(value: _protocols.ValueProtocol, *, name: str = "") -> onnx.
1818
1878
  return value_info_proto
1819
1879
 
1820
1880
 
1821
- @_capture_errors(lambda value_info_proto, from_: repr(from_))
1881
+ @_capture_errors(lambda value_info_proto, from_, name="": repr(from_))
1822
1882
  def serialize_value_into(
1823
1883
  value_info_proto: onnx.ValueInfoProto,
1824
1884
  from_: _protocols.ValueProtocol,
@@ -1,9 +1,9 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: onnx-ir
3
- Version: 0.1.6
3
+ Version: 0.1.7
4
4
  Summary: Efficient in-memory representation for ONNX
5
5
  Author-email: ONNX Contributors <onnx-technical-discuss@lists.lfaidata.foundation>
6
- License: Apache License v2.0
6
+ License-Expression: Apache-2.0
7
7
  Project-URL: Homepage, https://onnx.ai/ir-py
8
8
  Project-URL: Issues, https://github.com/onnx/ir-py/issues
9
9
  Project-URL: Repository, https://github.com/onnx/ir-py
@@ -13,7 +13,6 @@ Classifier: Programming Language :: Python :: 3.10
13
13
  Classifier: Programming Language :: Python :: 3.11
14
14
  Classifier: Programming Language :: Python :: 3.12
15
15
  Classifier: Programming Language :: Python :: 3.13
16
- Classifier: License :: OSI Approved :: Apache Software License
17
16
  Requires-Python: >=3.9
18
17
  Description-Content-Type: text/markdown
19
18
  License-File: LICENSE
@@ -1,56 +0,0 @@
1
- # Copyright (c) ONNX Project Contributors
2
- # SPDX-License-Identifier: Apache-2.0
3
- """Pass for removing duplicated initializer tensors from a graph."""
4
-
5
- from __future__ import annotations
6
-
7
- __all__ = [
8
- "DeduplicateInitializersPass",
9
- ]
10
-
11
-
12
- import onnx_ir as ir
13
-
14
-
15
- class DeduplicateInitializersPass(ir.passes.InPlacePass):
16
- """Remove duplicated initializer tensors from the graph.
17
-
18
- This pass detects initializers with identical shape, dtype, and content,
19
- and replaces all duplicate references with a canonical one.
20
-
21
- To deduplicate initializers from subgraphs, use :class:`~onnx_ir.passes.common.LiftSubgraphInitializersToMainGraphPass`
22
- to lift the initializers to the main graph first before running pass.
23
-
24
- .. versionadded:: 0.1.3
25
- """
26
-
27
- def __init__(self, size_limit: int = 1024):
28
- super().__init__()
29
- self.size_limit = size_limit
30
-
31
- def call(self, model: ir.Model) -> ir.passes.PassResult:
32
- graph = model.graph
33
- initializers: dict[tuple[ir.DataType, tuple[int, ...], bytes], ir.Value] = {}
34
- modified = False
35
-
36
- for initializer in tuple(graph.initializers.values()):
37
- # TODO(justinchuby): Handle subgraphs as well. For now users can lift initializers
38
- # out from the main graph before running this pass.
39
- const_val = initializer.const_value
40
- if const_val is None:
41
- # Skip if initializer has no constant value
42
- continue
43
-
44
- if const_val.size > self.size_limit:
45
- continue
46
-
47
- key = (const_val.dtype, tuple(const_val.shape), const_val.tobytes())
48
- if key in initializers:
49
- modified = True
50
- ir.convenience.replace_all_uses_with(initializer, initializers[key]) # type: ignore[index]
51
- assert initializer.name is not None
52
- graph.initializers.pop(initializer.name)
53
- else:
54
- initializers[key] = initializer # type: ignore[index]
55
-
56
- return ir.passes.PassResult(model=model, modified=modified)
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes