Skip to content
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 38 additions & 0 deletions src/onnx_ir/_convenience/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -556,3 +556,41 @@
)
value.type = new_value_type
return tensor


def all_values(
graph_like: _core.Graph | _core.Function | _core.GraphView,
) -> Iterable[_core.Value]:
"""Yield all values in a graph / function / graph view.

Args:
graph_like: The graph, function, or graph view to yield values from.

Returns:
An iterable of all values in the graph / function / graph view.
"""
# Yield all values in the model
seen = set()
for value in graph_like.inputs:
if value not in seen:
seen.add(value)
yield value
if not isinstance(graph_like, _core.Function):
for value in graph_like.initializers.values():
if value not in seen:
seen.add(value)
yield value
for node in traversal.RecursiveGraphIterator(graph_like):
# NOTE: If an input or initializer of a subgraph is unused, it will not be yielded by this function since it is not technically part of the graph. This is consistent with ONNX's handling of unused inputs and initializers.
for value in node.inputs:

Check failure

Code scanning / lintrunner

MYPY/assignment Error

Incompatible types in assignment (expression has type "Value | None", variable has type "Value") To disable, use # type: ignore[assignment]
if value not in seen:
seen.add(value)
yield value
for value in node.outputs:
if value not in seen:
seen.add(value)
yield value
for value in graph_like.outputs:
if value not in seen:
seen.add(value)
yield value
Loading