Decorator for trace_only functions

[justinchuby]

There is interest in enabling a trace_only function as a general ONNX Script feature. #674 (comment). Currently, we have the script: FunctionType -> OnnxFunction decorator that scripts a python function. Should we have a similar decorator for creating a trace only function?

Option 1

Retain script and create a new trace: FunctionType -> TracedOnnxFunction decorator. This is like torch.jit.trace and torch.jit.script, where the difference is we do symbolic tracing. TracedOnnxFunction is something that is OpLike and has scope information which we can use to dynamically generate functions during graph build.

Examples

@onnxscript.script()
def aten_abs(self: TReal) -> TReal:
    return op.Abs(self)

@onnxscript.trace()
def aten_abs(self: TReal) -> TReal:
    return op.Abs(self)

Option 2

Replace @script with @function, which does scripting by default, and add an option @function(trace=True). This is like tf.function with the jit_compile option (similar only on an api level. tf.function is always tracing I believe).

One of the pros of this option is a more natural representation of an ONNX "function".

Examples

@onnxscript.function()
def aten_abs(self: TReal) -> TReal:
    return op.Abs(self)

@onnxscript.function(trace=True)
def aten_abs(self: TReal) -> TReal:
    return op.Abs(self)

cc @gramalingam @BowenBao @titaiwangms @xadupre @jcwchen @fatcat-z @xiaowuhu @shubhambhokare1

[BowenBao]

Feels like we are only discussing api difference here?

Another key aspect to highlight is that in terms of graph creation. Is it correct to say script is aot compilation while trace is jit compilation, i.e. no onnx graph for trace decorated functions until feeding inputs (tensor can be maybe symbolic, but attributes must be concrete) to it?

[justinchuby]

Another key aspect to highlight is that in terms of graph creation. Is it correct to say script is aot compilation while trace is jit compilation, i.e. no onnx graph for trace decorated functions until feeding inputs (tensor can be maybe symbolic, but attributes must be concrete) to it?

That is how I understand it yes.

[titaiwangms]

Personally, I vote for option2. But I am not sure how much effort/cost/effect would be if we change the main api.