onnxscript.ir ¶

In-memory intermediate representation for ONNX graphs.

class onnxscript.ir.ArrayCompatible(*args, **kwargs)¶

Protocol for array-like objects.

An example of an array-like object is a numpy ndarray or a PyTorch Tensor. Read more at https://numpy.org/devdocs/user/basics.interoperability.html

class onnxscript.ir.Attr(name: str, type: AttributeType, value: Any, *, doc_string: str | None = None)¶

Base class for ONNX attributes.

as_float() → float[source]¶: Get the attribute value as a float.

as_floats() → Sequence[float][source]¶: Get the attribute value as a sequence of floats.

as_graph() → Graph[source]¶: Get the attribute value as a graph.

as_graphs() → Sequence[Graph][source]¶: Get the attribute value as a sequence of graphs.

as_int() → int[source]¶: Get the attribute value as an int.

as_ints() → Sequence[int][source]¶: Get the attribute value as a sequence of ints.

as_string() → str[source]¶: Get the attribute value as a string.

as_strings() → Sequence[str][source]¶: Get the attribute value as a sequence of strings.

as_tensor() → TensorProtocol[source]¶: Get the attribute value as a tensor.

as_tensors() → Sequence[TensorProtocol][source]¶: Get the attribute value as a sequence of tensors.

doc_string¶

name¶

type¶

value¶

onnxscript.ir.AttrFloat32(name: str, value: float, doc_string: str | None = None) → Attr¶: Create a float attribute.

onnxscript.ir.AttrFloat32s(name: str, value: Sequence[float], doc_string: str | None = None) → Attr¶: Create a float sequence attribute.

onnxscript.ir.AttrGraph(name: str, value: Graph, doc_string: str | None = None) → Attr¶: Create a graph attribute.

onnxscript.ir.AttrGraphs(name: str, value: Sequence[Graph], doc_string: str | None = None) → Attr¶: Create a graph sequence attribute.

onnxscript.ir.AttrInt64(name: str, value: int, doc_string: str | None = None) → Attr¶: Create an int attribute.

onnxscript.ir.AttrInt64s(name: str, value: Sequence[int], doc_string: str | None = None) → Attr¶: Create an int sequence attribute.

onnxscript.ir.AttrSparseTensor(name: str, value: SparseTensorProtocol, doc_string: str | None = None) → Attr¶: Create a sparse tensor attribute.

onnxscript.ir.AttrSparseTensors(name: str, value: Sequence[SparseTensorProtocol], doc_string: str | None = None) → Attr¶: Create a sparse tensor sequence attribute.

onnxscript.ir.AttrString(name: str, value: str, doc_string: str | None = None) → Attr¶: Create a str attribute.

onnxscript.ir.AttrStrings(name: str, value: Sequence[str], doc_string: str | None = None) → Attr¶: Create a string sequence attribute.

onnxscript.ir.AttrTensor(name: str, value: TensorProtocol, doc_string: str | None = None) → Attr¶: Create a tensor attribute.

onnxscript.ir.AttrTensors(name: str, value: Sequence[TensorProtocol], doc_string: str | None = None) → Attr¶: Create a tensor sequence attribute.

onnxscript.ir.AttrTypeProto(name: str, value: TypeAndShape, doc_string: str | None = None) → Attr¶: Create a type attribute.

onnxscript.ir.AttrTypeProtos(name: str, value: Sequence[TypeAndShape], doc_string: str | None = None) → Attr¶: Create a type sequence attribute.

class onnxscript.ir.AttributeProtocol(*args, **kwargs)¶

Protocol for ONNX attributes.

name¶

The name of the attribute.

Type:: str

type¶

The type of the attribute.

Type:: _enums.AttributeType

value¶

The value of the attribute.

Type:: Any

doc_string¶

Documentation string.

Type:: str | None

doc_string: str | None¶

name: str¶

type: _enums.AttributeType¶

value: Any¶

class onnxscript.ir.AttributeType(value)¶

Enum for the types of ONNX attributes.

FLOAT = 1¶

FLOATS = 6¶

GRAPH = 5¶

GRAPHS = 10¶

INT = 2¶

INTS = 7¶

SPARSE_TENSOR = 11¶

SPARSE_TENSORS = 12¶

STRING = 3¶

STRINGS = 8¶

TENSOR = 4¶

TENSORS = 9¶

TYPE_PROTO = 13¶

TYPE_PROTOS = 14¶

UNDEFINED = 0¶

class onnxscript.ir.DLPackCompatible(*args, **kwargs)¶

Protocol for objects that can support dlpack.

Computation backends can call __dlpack__ to obtain the underlying data in a tensor without copying the data. This allows use to use tensorflow tensors etc. without copying the data.

class onnxscript.ir.DataType(value)¶

Enum for the data types of ONNX tensors, defined in onnx.TensorProto.

BFLOAT16 = 16¶

BOOL = 9¶

COMPLEX128 = 15¶

COMPLEX64 = 14¶

DOUBLE = 11¶

FLOAT = 1¶

FLOAT16 = 10¶

FLOAT4E2M1 = 23¶

FLOAT8E4M3FN = 17¶

FLOAT8E4M3FNUZ = 18¶

FLOAT8E5M2 = 19¶

FLOAT8E5M2FNUZ = 20¶

INT16 = 5¶

INT32 = 6¶

INT4 = 22¶

INT64 = 7¶

INT8 = 3¶

STRING = 8¶

UINT16 = 4¶

UINT32 = 12¶

UINT4 = 21¶

UINT64 = 13¶

UINT8 = 2¶

UNDEFINED = 0¶

classmethod from_numpy(dtype: dtype) → DataType[source]¶

Returns the ONNX data type for the numpy dtype.

Raises:: TypeError – If the data type is not supported by ONNX.

property itemsize: float¶: Returns the size of the data type in bytes.

numpy() → dtype[source]¶

Returns the numpy dtype for the ONNX data type.

Raises:: TypeError – If the data type is not supported by numpy.

An immutable concrete tensor with its data store on disk.

This class uses memory mapping to avoid loading the tensor into memory, when the data type is supported by numpy. Otherwise, the tensor is loaded into memory lazily when accessed.

Calling shape does not incur IO. Checking shape before loading the tensor is recommended if IO overhead and memory usage is a concern.

To obtain an array, call numpy(). To obtain the bytes, call tobytes().

The location must be a relative path conforming to the ONNX specification. Given the correct base_dir, the path is computed to be the full path to the data file. Users should expect that the path always leads to the correct file. At initialization, paths are not checked. It is the user’s responsibility to ensure the paths are valid and accessible.

location¶: The location of the data file. It is the path relative to the base directory.

base_dir¶: The base directory for the external data. It is used to resolve relative paths. At serialization, only the location is serialized into the “location” field of the TensorProto.

path¶: The path to the data file. This is computed by joining base_dir and location.

offset¶: The offset in bytes from the start of the file.

length¶: The length of the data in bytes.

dtype¶: The data type of the tensor.

shape¶: The shape of the tensor.

name¶: The name of the tensor. It must be specified.

doc_string¶: The documentation string.

metadata_props¶: The metadata properties.

property base_dir: str | PathLike¶

doc_string¶

property dtype: DataType¶

property length: int | None¶

property location: str | PathLike¶

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]¶

name¶

numpy() → ndarray[source]¶

Return the tensor as a numpy array.

The data will be memory mapped into memory and will not taken up physical memory space.

property offset: int | None¶

property path: str¶

raw¶

release() → None[source]¶: Delete all references to the memory buffer and close the memory-mapped file.

property shape: Shape¶

tobytes() → bytes[source]¶

Return the bytes of the tensor.

This will load the tensor into memory.

class onnxscript.ir.Function(domain: str, name: str, overload: str = '', *, graph: Graph, attributes: Sequence[Attr], metadata_props: dict[str, str] | None = None)¶

IR functions.

Like a graph, a function can have nodes that are not topologically sorted. It is the responsibility of the user to maintain a topological order of the nodes.

Note that there is not a node attribute in the Function. The Function can be seen as a Sequence of nodes and should be used as such. For example, to obtain all nodes as a list, call list(function).

name¶: The function name.

domain¶: The domain this function is defined in.

overload¶: The overload name when the function is overloaded.

inputs¶: The input values of the function.

attributes¶: The attributes this function defines.

outputs¶: The output values of the function.

opset_imports¶: Opsets imported by the function.

doc_string¶: Documentation string.

metadata_props¶: Metadata that will be serialized to the ONNX file.

meta¶: Metadata store for graph transform passes.

append(node: Node, /) → None[source]¶: Append a node to the function in O(1) time.

property attributes: OrderedDict[str, Attr]¶

property doc_string: str | None¶

property domain: str¶

extend(nodes: Iterable[Node], /) → None[source]¶: Extend the function with the given nodes in O(#new_nodes) time.

identifier() → Tuple[str, str, str][source]¶: Return the unique identifier of the function.

property inputs: list[Value]¶

insert_after(node: Node, new_nodes: Iterable[Node], /) → None[source]¶: Insert new nodes after the given node in O(#new_nodes) time.

insert_before(node: Node, new_nodes: Iterable[Node], /) → None[source]¶: Insert new nodes before the given node in O(#new_nodes) time.

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]¶

property name: str¶

property opset_imports: dict[str, int]¶

property outputs: list[Value]¶

property overload: str¶

remove(nodes: Node | Iterable[Node], /, safe: bool = False) → None[source]¶

Remove nodes from the graph in O(#num of nodes) time.

If any errors are raise, to ensure the graph is not left in an inconsistent state, the graph is not modified.

Parameters:

nodes – The node to remove.
safe –
If True, performs the following actions before removal:

1. It checks to make sure there are no users of the node that are not to be removed before removing it. 2. It checks the node does not contribute to any graph outputs. 3. It removes references to all inputs so it is no longer a user of other nodes.

Raises:

ValueError – If any node to remove does not belong to this graph.
ValueError – (When safe=True) If the node does not belong to this graph or if there are users of the node.
ValueError – (When safe=True) If the node is still being used by other nodes not to be removed.

sort() → None[source]¶: Perform a topological sort of this graph and all subgraphs in O(#nodes + #values) time.

class onnxscript.ir.FunctionProtocol(*args, **kwargs)¶

Protocol for ONNX functions.

Like a graph, a function can have nodes that are not topologically sorted. It is the responsibility of the user to maintain a topological order of the nodes.

name¶

The function name.

Type:: str

domain¶

The domain this function is defined in.

Type:: str

overload¶

The overload name when the function is overloaded.

Type:: str

inputs¶

The input values of the function.

Type:: Sequence[ValueProtocol]

attributes¶

The attributes this function defines.

Type:: OrderedDict[str, AttributeProtocol]

outputs¶

The output values of the function.

Type:: Sequence[ValueProtocol]

opset_imports¶

Opsets imported by the function.

Type:: MutableMapping[str, int]

doc_string¶

Documentation string.

Type:: str

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

append(node: NodeProtocol, /) → None[source]¶: Append a node to the function.

attributes: OrderedDict[str, AttributeProtocol]¶

doc_string: str¶

domain: str¶

extend(nodes: Iterable[NodeProtocol], /) → None[source]¶: Extend the function with the given nodes.

identifier() → Tuple[str, str, str][source]¶: Return the unique identifier of the function.

inputs: Sequence[ValueProtocol]¶

insert_after(node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) → None[source]¶: Insert new nodes after the given node.

insert_before(node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) → None[source]¶: Insert new nodes before the given node.

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

name: str¶

opset_imports: MutableMapping[str, int]¶

outputs: Sequence[ValueProtocol]¶

overload: str¶

remove(node: NodeProtocol, /) → None[source]¶: Remove a node from the function.

sort() → None[source]¶: Topologically sort the nodes in the function.

class onnxscript.ir.Graph(inputs: Sequence[Value], outputs: Sequence[Value], *, nodes: Iterable[Node], initializers: Sequence[Value] = (), doc_string: str | None = None, opset_imports: dict[str, int] | None = None, name: str | None = None, metadata_props: dict[str, str] | None = None)¶

IR Graph.

Graph represents a computation graph. In addition to the ONNX specification specified fields, it also contains a mapping of opset_imports. This allows different subgraphs to import different opsets. It is the responsibility of the deserializer to reconcile the different opsets.

The nodes are not guaranteed to be topologically sorted. But the iteration order should be deterministic across different runs. It is the responsibility of the user to maintain a topological order of the nodes.

Note that there is not a node attribute in the Graph. The Graph can be seen as a Sequence of nodes and should be used as such. For example, to obtain all nodes as a list, call list(graph).

name¶: The name of the graph.

inputs¶: The input values of the graph.

outputs¶: The output values of the graph.

initializers¶: The initializers in the graph.

doc_string¶: Documentation string.

opset_imports¶: Opsets imported by the graph.

metadata_props¶: Metadata that will be serialized to the ONNX file.

meta¶: Metadata store for graph transform passes.

append(node: Node, /) → None[source]¶

Append a node to the graph in O(1) time.

Unique names will be assigned to the node and its values if any name is None.

Parameters:: node – The node to append.
Raises:: ValueError – If the node belongs to another graph.

property doc_string: str | None¶

extend(nodes: Iterable[Node], /) → None[source]¶

Extend the graph with the given nodes in O(#new_nodes) time.

Unique names will be assigned to the node and its values if any name is None.

Parameters:: nodes – The nodes to extend the graph with.
Raises:: ValueError – If any node belongs to another graph.

property initializers: dict[str, Value]¶

property inputs: list[Value]¶

insert_after(node: Node, new_nodes: Iterable[Node] | Node, /) → None[source]¶

Insert new nodes after the given node in O(#new_nodes) time.

Unique names will be assigned to the node and its values if any name is None.

Parameters:

node – The node to insert after.
new_nodes – The new nodes to insert.

Raises:

ValueError – If any node belongs to another graph.

insert_before(node: Node, new_nodes: Iterable[Node] | Node, /) → None[source]¶

Insert new nodes before the given node in O(#new_nodes) time.

Unique names will be assigned to the node and its values if any name is None.

Parameters:

node – The node to insert before.
new_nodes – The new nodes to insert.

Raises:

ValueError – If any node belongs to another graph.

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]¶

name¶

node(index_or_name: int | str, /) → Node[source]¶

Get a node by index or name.

This is an O(n) operation. Getting nodes on the ends of the graph (0 or -1) is O(1).

Note

If you need repeated random access, consider turning it into a list with list(graph) . Or a dictionary for repeated access by name: {node.name for node in graph} .

When a name is provided and if there are multiple nodes with the same name, the first node with the name is returned.

Parameters:

index_or_name – The index or name of the node.

Returns:

The node if found.

Raises:

IndexError – If the index is out of range.
ValueError – If the node with the given name is not found.

num_nodes() → int[source]¶

Get the number of nodes in the graph in O(1) time.

Note that this method returns the number of nodes this graph directly contains. It does not count nodes in subgraphs.

This is an alias for len(graph). Use this if you prefer a more descriptive name for readability.

property opset_imports: dict[str, int]¶

property outputs: list[Value]¶

register_initializer(value: Value) → None[source]¶

This is a convenience method to register an initializer to the graph with checks.

Parameters:

value – The Value to register as an initializer of the graph. It must have its .const_value set.

Raises:

ValueError – If a value of the same name that is not this value is already registered.
ValueError – If the value does not have a name.
ValueError – If the initializer is produced by a node.
ValueError – If the value does not have its .const_value set.

remove(nodes: Node | Iterable[Node], /, safe: bool = False) → None[source]¶

Remove nodes from the graph in O(#num of nodes to remove) time.

If any errors are raise, to ensure the graph is not left in an inconsistent state, the graph is not modified.

Parameters:

nodes – The node to remove.
safe –
If True, performs the following actions before removal:

1. It checks to make sure there are no users of the node that are not to be removed before removing it. 2. It checks the node does not contribute to any graph outputs. 3. It removes references to all inputs so it is no longer a user of other nodes.

Raises:

ValueError – If any node to remove does not belong to this graph.
ValueError – (When safe=True) If the node does not belong to this graph or if there are users of the node.
ValueError – (When safe=True) If the node is still being used by other nodes not to be removed.

sort() → None[source]¶

Perform a topological sort of this graph and all subgraphs in O(#nodes + #values) time.

This sort is stable. It preserves the original order as much as possible.

Referece: https://github.com/madelson/MedallionTopologicalSort#stable-sort

Raises:: ValueError – If the graph contains a cycle, making topological sorting impossible.

class onnxscript.ir.GraphProtocol(*args, **kwargs)¶

Protocol for graphs.

Note that there is not a node attribute in the Graph. The Graph can be seen as a Sequence of nodes and should be used as such. For example, to obtain all nodes as a list, call list(graph).

name¶

The name of the graph.

Type:: str | None

inputs¶

The input values of the graph.

Type:: MutableSequence[ValueProtocol]

outputs¶

The output values of the graph.

Type:: MutableSequence[ValueProtocol]

initializers¶

The initializers in the graph.

Type:: MutableMapping[str, ValueProtocol]

doc_string¶

Documentation string.

Type:: str

opset_imports¶

Opsets imported by the graph.

Type:: MutableMapping[str, int]

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

append(node: NodeProtocol, /) → None[source]¶: Append a node to the graph.

doc_string: str¶

extend(nodes: Iterable[NodeProtocol], /) → None[source]¶: Extend the graph with the given nodes.

initializers: MutableMapping[str, ValueProtocol]¶

inputs: MutableSequence[ValueProtocol]¶

insert_after(node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) → None[source]¶: Insert new nodes after the given node.

insert_before(node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) → None[source]¶: Insert new nodes before the given node.

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

name: str | None¶

opset_imports: MutableMapping[str, int]¶

outputs: MutableSequence[ValueProtocol]¶

remove(node: NodeProtocol, /) → None[source]¶: Remove a node from the graph.

sort() → None[source]¶: Topologically sort the nodes in the graph.

class onnxscript.ir.GraphView(inputs: Sequence[Value], outputs: Sequence[Value], *, nodes: Iterable[Node], initializers: Sequence[ValueProtocol] = (), doc_string: str | None = None, opset_imports: dict[str, int] | None = None, name: str | None = None, metadata_props: dict[str, str] | None = None)¶

A read-only view on a graph.

The GraphView is useful for analysis of a subgraph. It can be initialized with a subset of nodes from a Graph. Creating GraphView does not change the ownership of the nodes, and so it is possible to create multiple GraphViews that contain the same nodes. If the underlying nodes / connections are mutated, the mutation will be reflected in all views as well.

The graph view can be serialized to ONNX:

graph_proto = ir.serde.serialize_graph(graph_view)

It can also be used to create a model:

model = ir.Model(graph_view, ir_version=8)
model_proto = ir.serde.serialize_model(model)

The model created with a GraphView will have a fixed topology, and its graph will remain read-only as a GraphView. No copying will be done during the initialization process.

name¶: The name of the graph.

inputs¶: The input values of the graph.

outputs¶: The output values of the graph.

initializers¶: The initializers in the graph.

doc_string¶: Documentation string.

opset_imports¶: Opsets imported by the graph.

metadata_props¶: Metadata that will be serialized to the ONNX file.

meta¶: Metadata store for graph transform passes.

doc_string¶

initializers¶

inputs¶

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]¶

name¶

nodes¶

opset_imports¶

outputs¶

class onnxscript.ir.GraphViewProtocol(*args, **kwargs)¶

Protocol for a read-only view on a graph.

name¶

The name of the graph.

Type:: str | None

inputs¶

The input values of the graph.

Type:: Sequence[ValueProtocol]

outputs¶

The output values of the graph.

Type:: Sequence[ValueProtocol]

initializers¶

The initializers in the graph.

Type:: Mapping[str, ValueProtocol]

doc_string¶

Documentation string.

Type:: str

opset_imports¶

Opsets imported by the graph.

Type:: Mapping[str, int]

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

doc_string: str¶

initializers: Mapping[str, ValueProtocol]¶

inputs: Sequence[ValueProtocol]¶

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

name: str | None¶

opset_imports: Mapping[str, int]¶

outputs: Sequence[ValueProtocol]¶

onnxscript.ir.Input(name: str | None = None, shape: Shape | None = None, type: TypeProtocol | None = None, doc_string: str | None = None) → Value¶

Create an input of a Graph or a Function.

This is equivalent to calling Value(name=name, shape=shape, type=type, doc_string=doc_string).

class onnxscript.ir.MapTypeProtocol(*args, **kwargs)¶

Protocol for ONNX map types.

TODO: This protocol is not yet implemented in the ONNX IR.

key_type: Literal[_enums.DataType.STRING, _enums.DataType.INT64, _enums.DataType.INT32, _enums.DataType.INT16, _enums.DataType.INT8, _enums.DataType.UINT64, _enums.DataType.UINT32, _enums.DataType.UINT16, _enums.DataType.UINT8]¶

value_type: _enums.DataType¶

class onnxscript.ir.Model(graph: Graph, *, ir_version: int, producer_name: str | None = None, producer_version: str | None = None, domain: str | None = None, model_version: int | None = None, doc_string: str | None = None, functions: Sequence[Function] = (), meta_data_props: dict[str, str] | None = None)¶

doc_string¶

domain¶

property functions: dict[Tuple[str, str, str], Function]¶

graph¶

ir_version¶

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]¶

model_version¶

property opset_imports: dict[str, int]¶

producer_name¶

producer_version¶

class onnxscript.ir.ModelProtocol(*args, **kwargs)¶

Protocol for models.

A model is a container for a graph and metadata. It is the top-level object that represents an ONNX model.

graph¶

The graph of the model.

Type:: GraphProtocol

ir_version¶

The version of the IR.

Type:: int

producer_name¶

The name of the producer.

Type:: str | None

producer_version¶

The version of the producer.

Type:: str | None

domain¶

The domain of the model.

Type:: str | None

model_version¶

The version of the model.

Type:: int | None

doc_string¶

Documentation string.

Type:: str | None

functions¶

The functions defined in the model.

Type:: MutableMapping[str, FunctionProtocol]

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

doc_string: str | None¶

domain: str | None¶

functions: MutableMapping[str, FunctionProtocol]¶

graph: GraphProtocol¶

ir_version: int¶

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

model_version: int | None¶

opset_imports: MutableMapping[str, int]¶

producer_name: str | None¶

producer_version: str | None¶

IR Node.

If the graph is provided, the node will be added to the graph. Otherwise, user is responsible to call 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 the 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 replace the each of the inputs of output.uses() with the new output values by calling replace_input_with() on the using nodes of this node’s outputs.

append(nodes: Node | Iterable[Node]) → None[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 – A node or a sequence of nodes to put after this node.

property attributes: OrderedDict[str, Attr | RefAttr]¶

display(*, page: bool = False) → None[source]¶

Pretty print the object.

Parameters:: page – Whether to page the output.

doc_string¶

property domain: str¶

property graph: Graph | None¶

property inputs: Sequence[Value | None]¶

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]¶

property name: str | None¶

op_identifier() → Tuple[str, str, str][source]¶

property op_type: str¶

property outputs: Sequence[Value]¶

property overload: str¶

prepend(nodes: Node | Iterable[Node]) → None[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 – A node or a sequence of nodes to put before this node.

replace_input_with(index: int, value: Value | None) → None[source]¶: Replace an input with a new value.

property version: int | None¶

class onnxscript.ir.NodeProtocol(*args, **kwargs)¶

Protocol for nodes.

A node represents an invocation of an operation on the Value s in the computational graph.

A node can be optionally named. A name should typically be assigned when the node is added to a graph.

domain, op_type, and overload together uniquely identify the operator, and are always strings. For ONNX operators, domain and overload are both empty strings.

inputs and outputs are the input and output values of the node.

attributes are the attributes of the node. The attributes are stored in an ordered dictionary to preserve the order of the attributes. This is a deviation from the current ONNX spec where attributes are unordered, but it is helpful for tools that rely on the order of the attributes, e.g. those converting to and from Python function keyword arguments.

version is unique to the IR and is not specified in the ONNX spec. This allows the IR to represent a graph with mixed opset versions. Deserializers should decide how to reconcile the different versions within the graph. A typical graph will have a single version, declared in the Graph object and the nodes will have None as the version.

domain¶

The domain of the operator. E.g. "" for ONNX operators.

Type:: str

op_type¶

The operator name.

Type:: str

overload¶

The overload name when the node is invoking a function.

Type:: str

inputs¶

Input values.

Type:: Sequence[ValueProtocol]

outputs¶

Output values.

Type:: Sequence[ValueProtocol]

attributes¶

The attributes of the operator.

Type:: OrderedDict[str, AttributeProtocol | ReferenceAttributeProtocol]

version¶

The version of the operator.

Type:: int | None

doc_string¶

Documentation string.

Type:: str | None

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

attributes: OrderedDict[str, AttributeProtocol | ReferenceAttributeProtocol]¶

doc_string: str | None¶

domain: str¶

inputs: Sequence[ValueProtocol]¶

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

name: str | None¶

op_type: str¶

outputs: Sequence[ValueProtocol]¶

overload: str¶

replace_input_with(index: int, value: ValueProtocol | None) → None[source]¶: Set the input at the given index to the given value, replacing the original value.

version: int | None¶

class onnxscript.ir.OptionalType(elem_type: TypeProtocol, *, denotation: str | None = None)¶

A type that represents an optional element.

denotation: str | None¶

class onnxscript.ir.RefAttr(name: str, ref_attr_name: str, type: AttributeType, *, doc_string: str | None = None)¶

Reference attribute.

doc_string: str | None¶

property name: str¶

property ref_attr_name: str¶

property type: AttributeType¶

class onnxscript.ir.ReferenceAttributeProtocol(*args, **kwargs)¶

Protocol for a reference attribute.

A reference attribute can only appear inside the definition body of a function.

name¶

The name of the attribute.

Type:: str

ref_attr_name¶

The name of the attribute definition this attribute refers to.

Type:: str

type¶

The type of the attribute.

Type:: onnxscript.ir.AttributeType

doc_string¶

Documentation string.

Type:: str | None

doc_string: str | None¶

name: str¶

ref_attr_name: str¶

type: AttributeType¶

class onnxscript.ir.SequenceType(elem_type: TypeProtocol, *, denotation: str | None = None)¶

A type that represents a sequence of elements.

denotation: str | None¶

copy()[source]¶: Return a copy of the shape.

property dims: tuple[int | SymbolicDim, ...]¶

All dimensions in the shape.

This property is read-only. Use __getitem__ and __setitem__ to modify the shape or create a new shape.

get_denotation(index: int) → str | None[source]¶

Return the denotation of the dimension at the index.

Parameters:: index – The index of the dimension.
Returns:: The denotation of the dimension.

is_dynamic(dim=None) → bool[source]¶

is_static(dim=None) → bool[source]¶: Return True if the dimension is static. If dim is None, return True if all dimensions are static.

numpy() → tuple[int, ...][source]¶

rank() → int[source]¶: The rank of the shape.

set_denotation(index: int, denotation: str | None) → None[source]¶

Set the denotation of the dimension at the index.

Parameters:

index – The index of the dimension.
denotation – The denotation of the dimension.

class onnxscript.ir.ShapeProtocol(*args, **kwargs)¶

Protocol for ONNX shapes.

A shape is a sequence of dimensions.

dims¶

The dimensions of the shape.

Type:: Sequence[int | SymbolicDimProtocol]

dims: Sequence[int | SymbolicDimProtocol]¶

get_denotation(index: int) → str | None[source]¶

numpy() → Sequence[int][source]¶

rank() → int[source]¶

set_denotation(index: int, denotation: str | None) → None[source]¶

class onnxscript.ir.SparseTensorProtocol(*args, **kwargs)¶

dims: Sequence[int]¶

indices: TensorProtocol¶

values: TensorProtocol¶

class onnxscript.ir.SparseTensorType(dtype: DataType, *, denotation: str | None = None)¶

A type that represents a sparse tensor.

denotation: str | None¶

Multidimensional array of strings (as binary data to match the string_data field in TensorProto).

doc_string¶

property dtype: DataType¶: The data type of the tensor. Immutable.

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]¶

name¶

numpy() → npt.NDArray[np.bytes_][source]¶: Return the tensor as a numpy array.

property raw: Sequence[bytes] | npt.NDArray[np.bytes_]¶: Backing data of the tensor. Immutable.

property shape: Shape¶: The shape of the tensor. Immutable.

string_data() → Sequence[bytes][source]¶: Return the string data of the tensor.

tobytes() → bytes[source]¶: Return the tensor as a byte string conformed to the ONNX specification, in little endian.

class onnxscript.ir.SymbolicDim(value: str | None)¶

property value: str | None¶

class onnxscript.ir.SymbolicDimProtocol(*args, **kwargs)¶

Value of a single symbolic/dynamic dimension in a shape.

value¶

The value of the dimension.

Type:: str | None

value: str | None¶

An immutable concrete tensor.

This class is a wrapper around the raw tensor data. The raw tensor data can be a numpy array compatible object (e.g. np.ndarray, torch.Tensor) or a DLPack compatible object. The tensor is immutable and the data is not copied at initialization.

To create a tensor from a numpy array:

>>> import numpy as np
>>> array = np.array([1, 2, 3])
>>> tensor = Tensor(array)
>>> # The tensor itself can be treated as a numpy array because it implements the __array__ method
>>> np.allclose(tensor, array)
True

To get a numpy array from the tensor, call numpy(). To convert the tensor to a byte string for serialization, call tobytes().

It is recommended to check the size of the tensor first before accessing the underlying data, because accessing the data may be expensive and incur IO overhead.

Subclass this class to efficiently handle different types of tensors from different frameworks.

name¶: The name of the tensor.

shape¶: The shape of the tensor.

dtype¶: The data type of the elements of the tensor. It is an ir.DataType enum.

doc_string¶: Documentation string.

raw¶: The raw data behind this tensor. It can be anything.

size¶: The number of elements in the tensor.

nbytes¶: The number of bytes in the tensor.

metadata_props¶: Metadata that will be serialized to the ONNX file.

meta¶: Metadata store for graph transform passes.

doc_string¶

property dtype: DataType¶: The data type of the tensor. Immutable.

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]¶

name¶

numpy() → ndarray[source]¶

Return the tensor as a numpy array.

When the data type is not supported by numpy, the dtypes from the ml_dtype package are used. The values can be reinterpreted as bit representations using the .view() method.

property raw: TArrayCompatible¶: Backing data of the tensor. Immutable.

property shape: Shape¶: The shape of the tensor. Immutable.

tobytes() → bytes[source]¶

Returns the value as bytes encoded in little endian.

Override this method for more efficient serialization when the raw value is not a numpy array.

class onnxscript.ir.TensorProtoTensor(proto: TensorProto)¶

A tensor initialized from a tensor proto.

property doc_string: str¶

property dtype: DataType¶

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]¶

property name: str¶

numpy() → ndarray[source]¶

Return the tensor as a numpy array.

This is an improved version of onnx.numpy_helper.to_array. It first reads the data using the dtype corresponding to the tensor proto data field, then converts it to the correct dtype and shape. Special cases are bfloat16, complex and int4 where we need to reinterpret the data. Other types can simply be casted.

When the data type is not supported by numpy, the dtypes from the ml_dtype package are used. The values can be reinterpreted as bit representations using the .view() method.

When the data type is a string, this method returns a numpy array of bytes instead of a numpy array of strings, to follow the ONNX specification.

External tensors are not supported by this class. Use onnxscript.ir.ExternalTensor instead.

Raises:: ValueError – If the data type is UNDEFINED.

property raw: TensorProto¶

property shape: Shape¶

tobytes() → bytes[source]¶

Return the tensor as a byte string conformed to the ONNX specification, in little endian.

Raises:

ValueError – If the tensor is a string tensor or an external tensor.
ValueError – If the tensor is of UNDEFINED data type.

class onnxscript.ir.TensorProtocol(*args, **kwargs)¶

Concrete tensor backed by data.

The protocol does not specify how the data is stored. That data is exposed through the raw attribute for examination, but accessing raw is typically not needed.

To use the tensor as a numpy array, call numpy(). To convert the tensor to a byte string for serialization, call tobytes().

It is recommended to check the size of the tensor first before accessing the underlying data, because accessing the data may be expensive and incur IO overhead.

name¶

The name of the tensor.

Type:: str | None

shape¶

The shape of the tensor.

Type:: ShapeProtocol

dtype¶

The data type of the elements of the tensor. It is an ir.DataType enum.

Type:: _enums.DataType

doc_string¶

Documentation string.

Type:: str | None

raw¶

The raw data behind this tensor. It can be anything.

Type:: Any

size¶: The number of elements in the tensor.

nbytes¶: The number of bytes in the tensor.

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

doc_string: str | None¶

dtype: _enums.DataType¶

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

name: str | None¶

property nbytes: int¶

numpy() → np.ndarray[source]¶: Return the tensor as a numpy array.

raw: Any¶

shape: ShapeProtocol¶

property size: int¶

tobytes() → bytes[source]¶: Return the tensor as a byte string conformed to the ONNX specification, in little endian.

class onnxscript.ir.TensorType(dtype: DataType, *, denotation: str | None = None)¶

A type that represents a tensor.

denotation: str | None¶

class onnxscript.ir.TypeAndShape(type: TypeProtocol | None, shape: Shape | None)¶

Type and shape.

Useful for constructing a type proto.

shape: Shape | None¶

type: TypeProtocol | None¶

class onnxscript.ir.TypeProtocol(*args, **kwargs)¶

Protocol for ONNX tensors, Sequence tensors, Optional tensors and Sparse tensors.

These three types of tensors share the same attribute “elem_type” so they are merged in the same interface. Unlike the ONNX TensorProto, shapes are not included in the type and should be stored in the Value.

denotation¶

An optional denotation can be used to denote the whole type with a standard semantic description as to what is stored inside. Refer to https://github.com/onnx/onnx/blob/main/docs/TypeDenotation.md#type-denotation-definition for pre-defined type denotations.

Type:: str | None

elem_type¶

The type of its elements for nested types like Sequence[Optional] tensors. Or the DataType if the type is not nested.

Type:: onnxscript.ir.TypeProtocol | onnxscript.ir.DataType

dtype¶

The data type of the tensor or the nested tensor.

Type:: onnxscript.ir.DataType

denotation: str | None¶

dtype: DataType¶

elem_type: TypeProtocol | DataType¶

IR Value.

A value is a named entity that can be used to represent an input or output of a graph, a function, or a node. The information it stores generalizes over ValueInfoProto in the ONNX specification.

A Value is always not owned or owned by exactly one node. When the value is not owned, it must be an input of a graph or a function. producer and index are None.

When the value is owned by a node, it is an output of the node. The node that produces the value can be accessed with producer(). The index of the output of the node that produces the value can be accessed with index().

To find all the nodes that use this value as an input, call uses().

To check if the value is an output of a graph, call is_graph_output().

name¶: The name of the value. A value is always named when it is part of a graph.

shape¶: The shape of the value.

type¶: The type of the value.

metadata_props¶: Metadata.

property const_value: TensorProtocol | None¶

A concrete value.

The value can be backed by different raw data types, such as numpy arrays. The only guarantee is that it conforms TensorProtocol.

doc_string¶

property dtype: DataType | None¶: The data type of the tensor.

index() → int | None[source]¶: The index of the output of the defining node.

is_graph_output() → bool[source]¶: Whether the value is an output of a graph.

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]¶

property name: str | None¶

producer() → Node | None[source]¶

The node that produces this value.

When producer is None, the value does not belong to a node, and is typically a graph input or an initializer.

property shape: Shape | None¶

property type: TypeProtocol | None¶

The type of the tensor.

Example types can be TensorType, SparseTensorType, SequenceType, OptionalType. To obtain the data type of the tensor, use type.dtype or conveniently dtype.

uses() → Collection[tuple[Node, int]][source]¶

Return a set of uses of the value.

The set contains tuples of (Node, index) where the index is the index of the input of the node. For example, if node.inputs[1] == value, then the use is (node, 1).

class onnxscript.ir.ValueProtocol(*args, **kwargs)¶

Protocol for values.

A value is a named entity that can be used to represent an input or output of a graph, a function, or a node. The information it stores generalizes over ValueInfoProto in the ONNX specification.

A Value is always not owned or owned by exactly one node. When the value is not owned, it must be an input of a graph or a function. producer and index are None.

To find all the nodes that use this value as an input, call uses().

To check if the value is an output of a graph, call is_graph_output().

name¶

The name of the value. A value is always named when it is part of a graph.

Type:: str

shape¶

The shape of the value.

Type:: ShapeProtocol | None

type¶

The type of the value.

Type:: TypeProtocol | None

metadata_props¶

Metadata that will be serialized to the ONNX file.

Type:: MutableMapping[str, str]

meta¶

Metadata store for graph transform passes.

Type:: MutableMapping[str, Any]

doc_string¶

Documentation string.

Type:: str | None

const_value¶

The constant tensor is the value constant.

Type:: TensorProtocol | None

const_value: TensorProtocol | None¶

doc_string: str | None¶

index() → int | None[source]¶: The index of the output of the node that produces this value.

is_graph_output() → bool[source]¶: Whether this value is an output of a graph.

meta: MutableMapping[str, Any]¶

metadata_props: MutableMapping[str, str]¶

name: str¶

producer() → NodeProtocol | None[source]¶: The node that produces this value.

shape: ShapeProtocol | None¶

type: TypeProtocol | None¶

uses() → Collection[tuple[NodeProtocol, int]][source]¶: The set of (node, input_index) with node being those that use this value as an input.

onnxscript.ir.from_proto(proto: ModelProto | GraphProto | NodeProto | TensorProto | AttributeProto | ValueInfoProto | TypeProto | FunctionProto) → Any¶: Deserialize an ONNX proto message to an IR object.

onnxscript.ir.load(path: str | PathLike, format: str | None = None) → Model¶

Load an ONNX model from a file.

Parameters:

path – The path to the ONNX file.
format – The format of the file (e.g. protobuf, textproto, json, etc.). If None, the format is inferred from the file extension.

Returns:

The loaded model.

onnxscript.ir.save(model: Model, path: str | PathLike, format: str | None = None) → None¶

Save an ONNX model to a file.

Parameters:

model – The model to save.
path – The path to save the model to.
format – The format of the file (e.g. protobuf, textproto, json, etc.). If None, the format is inferred from the file extension.

Create a tensor value from an ArrayLike object or a TensorProto.

The dtype must match the value. Reinterpretation of the value is not supported, unless if the value is a plain Python object, in which case it is converted to a numpy array with the given dtype.

:param:`value` can be a numpy array, a plain Python object, or a TensorProto.

Example:

>>> from onnxscript import ir
>>> import numpy as np
>>> import ml_dtypes
>>> import onnx
>>> ir.tensor(np.array([1, 2, 3], dtype=np.int16))
Tensor<INT16,[3]>(array([1, 2, 3], dtype=int16), name=None)
>>> ir.tensor([1, 2, 3], dtype=ir.DataType.BFLOAT16)
Tensor<BFLOAT16,[3]>(array([1, 2, 3], dtype=bfloat16), name=None)
>>> tp_tensor = ir.tensor(onnx.helper.make_tensor("tensor", onnx.TensorProto.FLOAT, dims=[], vals=[0.5]))
>>> tp_tensor.numpy()
array(0.5, dtype=float32)

Parameters:

value – The numpy array to create the tensor from.
dtype – The data type of the tensor.
name – The name of the tensor.
doc_string – The documentation string of the tensor.

Returns:

A tensor value.

Raises:

ValueError – If the dtype does not match the value when value is not a plain Python object like list[int].

onnxscript.ir.to_proto(ir_object: ModelProtocol | GraphProtocol | NodeProtocol | ValueProtocol | AttributeProtocol | ReferenceAttributeProtocol | TensorProtocol | TypeProtocol | GraphViewProtocol | FunctionProtocol) → Any¶: Serialize an IR object to a proto.

onnxscript.ir¶

onnxscript.ir ¶