Node¶

class onnx_ir.Node(domain, op_type, inputs, attributes=(), *, overload='', num_outputs=None, outputs=None, version=None, graph=None, name=None, doc_string=None, metadata_props=None, device_configurations=())¶

IR Node.

Tip

For a more convenient way (that supports Python objects as attributes) to create a node, use the onnx_ir.node() constructor.

If graph is provided, the node will be added to the graph. Otherwise, the user is responsible for calling graph.append(node) (or other mutation methods in Graph) to add the node to the graph.

After the node is initialized, it will add itself as a user of its input values.

The output values of the node are created during node initialization and are immutable. To change the output values, create a new node and, for each use of the old outputs (output.uses()), replace the input in the consuming node by calling replace_input_with(). You can also use the replace_all_uses_with() method to replace all uses of the output values.

Note

When the domain is "ai.onnx", it is normalized to "".

Parameters:

domain (str)
op_type (str)
inputs (Iterable[Value | None])
attributes (Iterable[Attr] | Mapping[str, Attr])
overload (str)
num_outputs (int | None)
outputs (Sequence[Value] | None)
version (int | None)
graph (Graph | Function | None)
name (str | None)
doc_string (str | None)
metadata_props (dict[str, str] | None)
device_configurations (tuple[NodeDeviceConfiguration, ...])

device_configurations¶

doc_string¶

append(nodes)[source]¶

Insert a node after this node in the list of nodes in the graph.

It is the same as calling graph.insert_after(self, nodes).

Example:

Before: previous_node -> self
        previous_node' -> node -> next_node'
After:  previous_node -> self -> node
        previous_node' -> next_node'

Parameters:: nodes (Node | Iterable[Node]) – A node or a sequence of nodes to put after this node.
Return type:: None

property attributes: Attributes¶

The attributes of the node as dict[str, Attr] with additional access methods.

Use it as a dictionary with keys being the attribute names and values being the Attr objects.

Use node.attributes.add(attr) to add an attribute to the node. Use node.attributes.get_int(name, default) to get an integer attribute value. Refer to the Attributes for more methods.

display(*, page=False)[source]¶

Pretty print the node.

This method is used for debugging and visualization purposes.

Parameters:: page (bool)
Return type:: None

property domain: str¶: The domain of the operator. For onnx operators, this is an empty string.

property graph: Graph | None¶

The graph that the node belongs to.

If the node is not added to any graph, this property is None.

property inputs: Sequence[Value | None]¶

The input values of the node.

The inputs are immutable. To change the inputs, create a new node and replace the inputs of the using nodes of this node’s outputs by calling replace_input_with() on the using nodes of this node’s outputs.

property meta: MetadataStore¶

The metadata store for intermediate analysis.

Write to the metadata_props if you would like the metadata to be serialized to the ONNX proto.

property metadata_props: dict[str, str]¶

The metadata properties of the node.

The metadata properties are used to store additional information about the node. Unlike meta, this property is serialized to the ONNX proto.

property name: str | None¶: Optional name of the node.

op_identifier()[source]¶

Return the operator identifier of the node.

The operator identifier is a tuple of the domain, op_type and overload.

Return type:: tuple[str, str, str]

property op_type: str¶: The name of the operator called.

property outputs: Sequence[Value]¶

The output values of the node.

The outputs are always attached to this node once initialized (immutable), except that the list can be resized to remove or add outputs.

Use resize_outputs() to change the number of outputs of the node.

property overload: str¶: The overload name when the node is invoking a function.

predecessors()[source]¶

Return the predecessor nodes of the node, deduplicated, in a deterministic order.

Return type:: Sequence[Node]

prepend(nodes)[source]¶

Insert a node before this node in the list of nodes in the graph.

It is the same as calling graph.insert_before(self, nodes).

Example:

Before: previous_node -> self
        previous_node' -> node -> next_node'
After:  previous_node -> node -> self
        previous_node' -> next_node'

Parameters:: nodes (Node | Iterable[Node]) – A node or a sequence of nodes to put before this node.
Return type:: None

replace_input_with(index, value)[source]¶

Replace an input with a new value.

Parameters:

index (int)
value (Value | None)

Return type:

None

resize_inputs(new_size, /)[source]¶

Resize the inputs of the node.

If the new size is greater than the current size, new inputs are added as None. If the new size is less than the current size, the extra inputs are removed.

After inputs is resized, you can use replace_input_with() to set the new inputs.

Added in version 0.1.13.

Parameters:: new_size (int) – The new number of inputs.
Return type:: None

resize_outputs(new_size, /)[source]¶

Resize the outputs of the node.

If the new size is greater than the current size, new output values are created. If the new size is less than the current size, the extra output values are removed. The removed output values must not have any uses.

Added in version 0.1.13.

Parameters:: new_size (int) – The new number of outputs.
Raises:: ValueError – If the new size is less than the current size and the removed outputs have uses.
Return type:: None

set_pipeline_stage(configuration, stage)[source]¶

Assign this node to a pipeline stage for configuration.

This expresses pure device placement / pipeline parallelism — for example putting a contiguous block of decoder layers on one device. Unlike shard(), it attaches no sharding spec; the node is placed as a whole.

If the node already has a NodeDeviceConfiguration for configuration (for example one created by shard()), its pipeline_stage is updated in place; otherwise a new placement-only configuration is created. Calling this method again replaces the stage.

How a stage maps to a physical device is by convention: a common choice is stage == device index into configuration.device_names.

Added in version 1.0.0.

Parameters:

configuration (ModelConfiguration) – The ModelConfiguration the stage belongs to.
stage (int) – The pipeline stage index. Must be >= 0.

Raises:

ValueError – If stage is negative.

Return type:

None

shard(value, *, configuration, axis, num_shards, device_indices=(), pipeline_stage=None)[source]¶

Record a sharding of value along axis for this node.

This is a convenience wrapper that builds and attaches the appropriate multi-device configuration metadata in device_configurations. The sharding is bound to value and configuration by object identity, so it follows renames and survives as long as the objects do.

Calling this method repeatedly for the same value and configuration but different axes builds a single ShardingSpec with one ShardedDim per axis (the canonical representation for sharding a tensor across a multi-axis device mesh). device_indices are unioned across those calls.

Added in version 1.0.0.

Parameters:

value (Value) – The input or output value to shard. Must be one of this node’s own inputs or outputs.
configuration (ModelConfiguration) – The ModelConfiguration the sharding belongs to.
axis (int) – The axis along which value is sharded. May be negative to count from the back (in the range [-rank, rank) following the ONNX convention), when the rank of value is known.
num_shards (int) – The number of shards along axis. Must be >= 1.
device_indices (Sequence[int]) – Optional indices (into configuration.device_names) of the devices the tensor is placed on. These are device indices, not the device names passed to Model.add_device_configuration().
pipeline_stage (int | None) – Optional pipeline stage for the configuration. Must be >= 0 and must not conflict with a stage already set on the configuration.

Raises:

ValueError – If value is not an input or output of the node, if num_shards < 1, if pipeline_stage is negative, if axis is out of range when the rank of value is known, if value is already sharded along axis for this configuration, or if pipeline_stage conflicts with the configuration’s existing stage.

Return type:

None

sharding_of(value)[source]¶

Return all sharding specs on this node that target value.

Matching is by object identity, so this returns the live specs that reference exactly value.

Added in version 1.0.0.

Parameters:: value (Value)
Return type:: tuple[ShardingSpec, …]

successors()[source]¶

Return the successor nodes of the node, deduplicated, in a deterministic order.

Return type:: Sequence[Node]

property version: int | None¶

Opset version of the operator called.

If None, the version is unspecified and will follow that of the graph. This property is special to ONNX IR to allow mixed opset usage in a graph for supporting more flexible graph transformations. It does not exist in the ONNX serialization (protobuf) spec.