Custom pseudo-operations for generic rewriting

[leshabirukov]

This text organized as

1. Problem
2. Solution
3. This issue subject

Problem to solve

I want element-wise operations move above their commutatives (friends). See picture:

I don't want to encode every possible interchange pair, so must be generic element-wise and generic friend. Can I have it with current onnxscript (0.2.2) rewriter?

2. Solution

Use meta-abc node pairs like this:

In the attachment, there is prototype implementation.
It moves elementwise operations (ReLU and Tanh) above Identity and Flatten

cusops2ons.py.txt

test_generic_as_wrap.py.txt

For now, we encode elementwise wrappers as [Split(axis=20), Sub]
and wrappers of their friend as [Split(axis=30), Sub]

What do you think about this approach to generics?

3. The issue

Well, it works! But using Split and Sub as replacement to ABC_pre and ABC_post is obviously bad.
I actually tried to create custom pseudo-operation, but failed.

As I create op like this:

my_opset = Opset("com.mydomain", 15)

@script(my_opset)
def joy_of_relu( x:FLOAT ) -> FLOAT:
    y =    op.Relu(x)
    return op.Relu(y)


opsch = joy_of_relu.op_schema

shiny_new_op = Op(
    opset=my_opset, name="relurelu", op_schema=opsch
)

it not work with rewriter, also I tried onnxruntime_extensions.onnx_op

So, if it is possible, it would be great to have example of using rewriter with custom nodes.