Define a custom op with a custom name

[Johansmm]

I would like to know if it is possible to define a custom operator with a different name that its function name.

Example:

import onnx

# We use ONNX opset 15 to define the function below.
from onnxscript import FLOAT, script
from onnxscript import opset15 as op


@script(op_type="MY_NEW_NAME_OP")
def custom_op(X):
    return 5*X


def sample_model(X: FLOAT[64, 128]) -> FLOAT[64, 128]:
    return custom_op(X)


# onnx_model is an in-memory ModelProto
onnx_model = sample_model.to_model_proto()

# Check node op_type
assert model.graph.node[0].op_type == "MY_NEW_NAME_OP"

[justinchuby]

Could you share your usage? When would this be useful?

[Johansmm]

My use case can be explained with the gelu function. As a rule, function names in python are usually written in snake or camel cases, but ONNX legacy operators are capital case. That is why if I defined gelu function with onnscript :

@script()
def gelu(X):
    ...

gelu.to_function_proto() method should map a ONNX function named 'gelu' and not 'Gelu'.

It is in this kind of cases that I think it is convenient to have the option to serialize the operator with a different name than the python function.

[justinchuby]

Makes sense. Thanks! cc @gramalingam