torch.onnx.ops

ONNX operators as native torch.fx operators.

This module provides a set of functions to create ONNX operators in the FX graph which are exportable to ONNX.

Operators

torch.onnx.ops.symbolic(domain_op, /, inputs, attrs=None, *, dtype, shape, version=None, metadata_props=None)

Create a symbolic FX operator to represent an arbitrary ONNX operator.

This function is used to create a symbolic operator with a single output. To create an operator with multiple outputs, use symbolic_multi_out().

Example:

class CustomOp(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Normal torch operators can interleave with the symbolic ops during ONNX export
        x = x + 1

        # Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
        # The output tensor will have the specified dtype and shape
        val = torch.onnx.ops.symbolic(
            "custom_domain::CustomOp",
            (x,),
            dict(attr_key="attr_value"),
            dtype=x.dtype,
            shape=x.shape,
            version=1,
        )

        # The result of the symbolic op can be used in normal torch operations during ONNX export
        return torch.nn.functional.relu(val)


# You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).

Parameters

• domain_op (str) – The domain and operator name, separated by “::”. For example, “custom_domain::CustomOp”.

• inputs (Sequence[torch.Tensor | None]) – The input tensors to the operator.

• attrs (dict[str, int | float | str | bool | Sequence[int] | Sequence[float] | Sequence[str] | Sequence[bool]] | None) – The attributes of the operator. The keys are attribute names and the values are attribute values. Valid attribute types are int, float, str, bool, and lists of int, float, str, and bool. Tensor attributes are unsupported.

• dtype (torch.dtype | int) – The data type of the output tensor.This can be either a torch.dtype or an integer representing the ONNX data type.

• shape (Sequence[int | torch.SymInt]) – The shape of the output tensor. This can be a list of integers or SymInt values.

• version (int | None) – The version of the opset used for the operator.

• metadata_props (dict[str, str] | None) – Metadata properties for the ONNX node. This is a dictionary of str-str pairs.

Returns

The output tensor of the operator.

Return type

torch.Tensor

torch.onnx.ops.symbolic_multi_out(domain_op, /, inputs, attrs=None, *, dtypes, shapes, version=None, metadata_props=None)

Create a symbolic FX operator to represent an arbitrary ONNX operator with multiple outputs.

Example:

class CustomOp(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Normal torch operators can interleave with the symbolic ops during ONNX export
        x = x + 1

        # Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
        # The output tensors will have the specified dtypes and shapes
        (out1, out2) = torch.onnx.ops.symbolic(
            "custom_domain::CustomOp",
            (x,),
            dict(attr_key="attr_value"),
            dtypes=(x.dtype, torch.float32),
            shapes=(x.shape, [1, 2, 3]),
            version=1,
        )

        # The result of the symbolic op can be used in normal torch operations during ONNX export
        return torch.nn.functional.relu(out1 + out2)


# You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).

Parameters

• domain_op (str) – The domain and operator name, separated by “::”. For example, “custom_domain::CustomOp”.

• inputs (Sequence[torch.Tensor | None]) – The input tensors to the operator.

• attrs (dict[str, int | float | str | bool | Sequence[int] | Sequence[float] | Sequence[str] | Sequence[bool]] | None) – The attributes of the operator. The keys are attribute names and the values are attribute values. Valid attribute types are int, float, str, bool, and lists of int, float, str, and bool. Tensor attributes are unsupported.

• dtypes (Sequence[torch.dtype | int]) – The data types of the output tensors. This can be a list of torch.dtype or integers representing the ONNX data types. The length of this list must be the number of outputs.

• shapes (Sequence[Sequence[int | torch.SymInt]]) – The shapes of the output tensors. This can be a list of lists of integers or SymInt values. The length of this list must be the number of outputs.

• version (int | None) – The version of the opset used for the operator.

• metadata_props (dict[str, str] | None) – Metadata properties for the ONNX node. This is a dictionary of str-str pairs.

Returns

A list of output tensors of the operator.

Return type

Sequence[torch.Tensor]