102: Examples with onnxscript

This script gathers a couple of examples based on onnxscript.

Custom Opset and Local Functions

import onnx
import onnxscript

op = onnxscript.opset18
my_opset = onnxscript.values.Opset("m_opset.ml", version=1)


@onnxscript.script(my_opset, default_opset=op)
def do_this(x, y):
    return op.Add(x, y)


@onnxscript.script(my_opset, default_opset=op)
def do_that(x, y):
    return op.Sub(x, y)


@onnxscript.script(my_opset, default_opset=op)
def do_this_or_do_that(x, y, do_this_or_do_that: bool = True):
    if do_this_or_do_that:
        ret = my_opset.do_this(x, y)
    else:
        ret = my_opset.do_that(x, y)
    return ret

/home/xadupre/github/onnxscript/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()
'do_this' is not a known op in 'm_opset.ml'
'do_that' is not a known op in 'm_opset.ml'

Then we export the model into ONNX.

proto = do_this_or_do_that.to_model_proto(functions=[do_this, do_that])
print(onnx.printer.to_text(proto))

<
   ir_version: 8,
   opset_import: ["" : 18, "m_opset.ml" : 1]
>
do_this_or_do_that ( x,  y) => ( ret_1) {
   [n0] do_this_or_do_that = Constant <value_int: int = @do_this_or_do_that> ()
   [n1] do_this_or_do_that_as_bool = Cast <to: int = 9> (do_this_or_do_that)
   [n2] ret_1 = If (do_this_or_do_that_as_bool) <then_branch: graph = thenGraph_3 () => ( ret) {
      [n0] ret = m_opset.ml.do_this (x, y)
   }, else_branch: graph = elseGraph_3 () => ( ret_0) {
      [n0] ret_0 = m_opset.ml.do_that (x, y)
   }>
}
<
  domain: "m_opset.ml",
  opset_import: ["" : 18]
>
do_this (x, y) => (return_val)
{
   [n0] return_val = Add (x, y)
}
<
  domain: "m_opset.ml",
  opset_import: ["" : 18]
>
do_that (x, y) => (return_val)
{
   [n0] return_val = Sub (x, y)
}