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.

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(X <= zero, 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

  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 = LessOrEqual (X, zero)
   return_val = Where (tmp_2, neg, pos)

Total running time of the script: (0 minutes 0.014 seconds)

Gallery generated by Sphinx-Gallery