to_onnx and a custom operator inplace

This example shows how to convert a custom operator as defined in the tutorial Python Custom Operators.

Inplace modification are not supported by onnx.

A model with a custom ops

from typing import Any, Dict, List, Optional
import numpy as np
import torch
from onnx_array_api.plotting.graphviz_helper import plot_dot
from experimental_experiment.xbuilder import GraphBuilder
from experimental_experiment.helpers import pretty_onnx
from experimental_experiment.torch_interpreter import to_onnx, Dispatcher

We define a model with a custom operator.

@torch.library.custom_op("mylib::numpy_sin", mutates_args={"output"}, device_types="cpu")
def numpy_sin(x: torch.Tensor, output: torch.Tensor) -> None:
    assert x.device == output.device
    assert x.device.type == "cpu"
    x_np = x.numpy()
    output_np = output.numpy()
    np.sin(x_np, out=output_np)


class ModuleWithACustomOperator(torch.nn.Module):
    def forward(self, x):
        out = torch.zeros(x.shape)
        numpy_sin(x, out)
        return out


model = ModuleWithACustomOperator()

Let’s check it runs.

x = torch.randn(1, 3)
model(x)

tensor([[0.2777, 0.5496, 0.7560]])

As expected, it does not export.

try:
    torch.export.export(model, (x,))
    raise AssertionError("This export should failed unless pytorch now supports this model.")
except Exception as e:
    print(e)

This export should failed unless pytorch now supports this model.

The exporter fails with the same eror as it expects torch.export.export to work.

try:
    to_onnx(model, (x,))
except Exception as e:
    print(e)

Unable to interpret function <class 'torch._ops.OpOverload'>: <OpOverload(op='mylib.numpy_sin', overload='default')>, searched for ['mylib::numpy_sin', 'numpy_sin_default'] and attributes ['__qualname__', '__name__'], args=(<OpOverload(op='mylib.numpy_sin', overload='default')>,), kwargs={'x': x, 'output': zeros}
--DEBUG--
[GraphBuilder-WKU] Message starts, there are 1 initializers, 1 nodes, 1 inputs, 1 outputs.
--LOCAL FUNCTIONS--
--PARAMETERS--
dynamic_examples=
--SHAPE--
dynamic_examples=
dynamic_objects=
dynamic_objects_rev=
dynamic_dimensions_source={}
dynamic_alias={}
dynamic_shapes=None
_known_value_shape={}
_known_types={'init7_s2_1_3': 7, 'x': 1, 'zeros': 1}
_known_shapes={'init7_s2_1_3': (2,), 'x': (1, 3), 'zeros': (1, 3)}
_known_constants=['init7_s2_1_3', 'zeros']
_known_ranks={}
--TORCH-USERS--
auto_functionalized -> {getitem_1}
x -> {auto_functionalized}
zeros -> {auto_functionalized}
--TORCH-SHAPES--
x: ('run_node', ('', ('val', torch.float32, torch.Size([1, 3])))) --- 1:2:(1, 3):
zeros: ('run_node', ('', ('val', torch.float32, torch.Size([1, 3])))) --- 1:2:(1, 3):
auto_functionalized: ('run_node', ('', '')) --- :::
--ONNX--
-- process.graph_module --
graph():
    %x : [num_users=1] = placeholder[target=x]
    %zeros : [num_users=1] = call_function[target=torch.ops.aten.zeros.default](args = ([1, 3],), kwargs = {device: cpu, pin_memory: False})
    %auto_functionalized : [num_users=1] = call_function[target=torch.ops.higher_order.auto_functionalized](args = (mylib.numpy_sin.default,), kwargs = {x: %x, output: %zeros})
    %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%auto_functionalized, 1), kwargs = {})
    return (getitem_1,)
-- process.progress --
node 2/5 target=auto_functionalized
--
[GraphBuilder-WKU.make_tensor_input] x[1:1x3]
[GraphBuilder-WKU.make_initializer] init7_s2_1_3[int64:int64:[1, 3]]
[GraphBuilder-WKU.make_node] zeros           [#:#   ] ConstantOfShape:['init7_s2_1_3']->['zeros']
[GraphBuilder-WKU] Message completed, there are 1 initializers, 1 nodes, 1 inputs, 1 outputs.

Registration

The exporter how to convert the new exporter into ONNX. This must be defined. The first piece is to tell the exporter that the shape of the output is the same as x. input names must be the same.

@numpy_sin.register_fake
def numpy_sin_shape(x, output):
    pass

Let’s see what the fx graph looks like.

print(torch.export.export(model, (x,)).graph)

graph():
    %x : [num_users=1] = placeholder[target=x]
    %zeros : [num_users=1] = call_function[target=torch.ops.aten.zeros.default](args = ([1, 3],), kwargs = {device: cpu, pin_memory: False})
    %auto_functionalized : [num_users=1] = call_function[target=torch.ops.higher_order.auto_functionalized](args = (mylib.numpy_sin.default,), kwargs = {x: %x, output: %zeros})
    %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%auto_functionalized, 1), kwargs = {})
    return (getitem_1,)

Next is the conversion to onnx.

T = str  # a tensor name


def numpy_sin_to_onnx(
    g: GraphBuilder,
    sts: Dict[str, Any],
    outputs: List[str],
    x: T,
    output: Optional[T] = None,
    name: str = "mylib.numpy_sin",
) -> T:
    # name= ... lets the user know when the node comes from
    # o is not used, we could check the shape are equal.
    # outputs contains unexpectedly two outputs
    g.op.Sin(x, name=name, outputs=outputs[1:])
    return outputs

We create a Dispatcher.

dispatcher = Dispatcher({"mylib::numpy_sin": numpy_sin_to_onnx})

And we convert again.

onx = to_onnx(model, (x,), dispatcher=dispatcher, optimize=False)
print(pretty_onnx(onx))

opset: domain='' version=18
doc_string: large_model=False, inline=False, external_threshold=102...
input: name='x' type=dtype('float32') shape=[1, 3]
init: name='init7_s2_1_3' type=dtype('int64') shape=(2,) -- array([1, 3])
ConstantOfShape(init7_s2_1_3, value=[nan]) -> zeros
Sin(x) -> auto_functionalized#1
  Identity(auto_functionalized#1) -> getitem_1
    Identity(getitem_1) -> output_0
output: name='output_0' type=dtype('float32') shape=[1, 3]

And we convert again with optimization this time.

onx = to_onnx(model, (x,), dispatcher=dispatcher, optimize=True)
print(pretty_onnx(onx))

opset: domain='' version=18
doc_string: large_model=False, inline=False, external_threshold=102...
input: name='x' type=dtype('float32') shape=[1, 3]
Sin(x) -> output_0
output: name='output_0' type=dtype('float32') shape=[1, 3]

And visually.

plot_dot(onx)

<Axes: >