Generating a FunctionProto

The example below shows how we can define Selu as a function in onnxscript.

First, import the ONNX opset used to define the function.

from onnxscript import opset15 as op
from onnxscript import script

Next, define Selu as an ONNXScript function.

@script()
def Selu(X, alpha: float, gamma: float):
    alphaX = op.CastLike(alpha, X)
    gammaX = op.CastLike(gamma, X)
    neg = gammaX * (alphaX * op.Exp(X) - alphaX)
    pos = gammaX * X
    zero = op.CastLike(0, X)
    return op.Where(zero >= X, neg, pos)

We can convert the ONNXScript function to an ONNX function (FunctionProto) as below:

onnx_fun = Selu.to_function_proto()

Let’s see what the translated function looks like:

import onnx  # noqa: E402

print(onnx.printer.to_text(onnx_fun))

<
  domain: "this",
  opset_import: ["" : 15]
>
Selu <alpha,gamma>(X) => (return_val)
{
   alpha = Constant <value_float: float = @alpha> ()
   alphaX = CastLike (alpha, X)
   gamma = Constant <value_float: float = @gamma> ()
   gammaX = CastLike (gamma, X)
   tmp = Exp (X)
   tmp_0 = Mul (alphaX, tmp)
   tmp_1 = Sub (tmp_0, alphaX)
   neg = Mul (gammaX, tmp_1)
   pos = Mul (gammaX, X)
   int64_0 = Constant <value: tensor = int64 int64_0 {0}> ()
   zero = CastLike (int64_0, X)
   tmp_2 = GreaterOrEqual (zero, X)
   return_val = Where (tmp_2, neg, pos)
}