onnx_extended.tools.graph

NodeKind

class onnx_extended.tools.graph.onnx_graph_struct.NodeKind(value)

Node kind.

Node

class onnx_extended.tools.graph.onnx_graph_struct.Node(index: int, parent: Graph, proto: TensorProto | NodeProto | ValueInfoProto | str, kind: NodeKind | None = None)

Defines a node in the graph. It can be an iniatialier or a node.

create_with_new_values(new_tensor: TensorProto) → Node

Creates an iniatializer or a node Constant based on the new value.

get_tensor() → TensorProto

Returns the value of the

getattr(name: str, astype: type | None = None, has_default: bool = False) → Any

Retrieves a specific attribute and extracts its value if astype is not None.

Parameters:

• name – attribute name

• astype – cast the attribute into this type

• has_default – if the parameter has a default value, the method returns None if the attribute is not found

Returns:

the value of the attribute or an AttributeProto if astype is None

property inputs: List[str]

Input names

is_constant() → bool

True if operator Constant or initializer or a Constant as an output of an operator taking only constants.

property is_initializer: bool

True if inititializer

property is_input: bool

True if an input

property is_node: bool

True if a NodeProto.

property is_output: bool

True if an output

property is_sparse_initializer: bool

True if inititializer

match(pattern: Dict[str, str] | None) → bool

Checks if a node match the proposed pattern.

Parameters:

pattern – a node matches the pattern {“name”: “node_name”} if its node is equal to ‘node_name’

Returns:

match

property name

Returns the name if node is a NodeProto, None otherwise.

property op_type: str

Returns the node type.

property outname

Returns the output name.

property outputs: List[str]

Output names

NodeWithSubGraph

class onnx_extended.tools.graph.onnx_graph_struct.NodeWithSubGraph(index: int, parent: Graph, proto: NodeProto)

A node with a subgraphs (If, Loop, Scan, …).

property inputs: List[str]

Input names

NodeSet

class onnx_extended.tools.graph.onnx_graph_struct.NodeSet(nodes: List[Node])

Defines a set of nodes.

Graph

class onnx_extended.tools.graph.Graph(proto: FunctionProto | GraphProto | ModelProto)

A GraphProto, FunctionProto or ModelProto.

add_functions(protos: Iterable[FunctionProto])

Adds functions to the graph when it is exported to ONNX.

Parameters:

protos – enumerate of FunctionProto

generate_name(prefix: str = 'new') → str

Generates a name which is not used for any existing result in the graph.

Parameters:

prefix – prefix to use for the new name, next tries will be <prefix>_1, <prefix>_2, …

Returns:

new name

generate_node_name(prefix: str = 'new') → str

Generates a node name which is not used for any existing node in the graph.

Parameters:

prefix – prefix to use for the new name, next tries will be <prefix>_1, <prefix>_2, …

Returns:

new name

get_node_producer(name: str) → Node

Returns the node producing the output name.

Parameters:

name – output name to check

Returns:

Node producing the output name or None if it is an input.

get_opset(domain: str = '') → int

Returns the opset for a specific domain.

Parameters:

domain – domain

Returns:

model opset

get_opsets() → Dict[str, int]

Returns the opsets available registered for ever domain in the model.

get_shape(name: str) → Tuple[int | str | None, ...] | None

Returns the shape of a result.

Parameters:

name – name of the result

Returns:

None if unknown or a tuple

is_constant(name: str) → bool

Tells if output name is constant.

Parameters:

name – result name

Returns:

True if constant

remove_unused_nodes()

Removes unused nodes, a node with only unused outputs.

Returns:

removed nodes

replace_nodes(indices: int | List[int], new_nodes: NodeProto | List[NodeProto], new_opsets: Dict[str, int] | None = None) → NodeSet

Replaces a node index

Parameters:

• indices – index or list of indices to replace

• new_nodes – node or list of nodes to add

• new_opsets – new opet versions

Returns:

added nodes

By default, the nodes are inserted at position indices[-1]. It ensures the inputs of the new nodes were already computed. However, it does not ensure that every intermediate node between the first and the last removed nodes can be computed. Sorting the nodes is needed in that. This function does not do that.

simplify(remove_unused: bool = True) → Graph

Stores every node into nodes. Removes unused nodes.

Parameters:

remove_unused – removes unused nodes as well, see remove_unused_nodes()

Returns:

self

to_onnx() → ModelProto | FunctionProto | GraphProto

Converts the current graph into onnx with the same type as the input type.

upgrade_opsets(new_opsets: Dict[str, int])

Upgrades the models to a newer opsets.

Parameters:

new_opsets – dictionary { domain: new version }