megengine.traced_module.wrap#

wrap(func)[source]#

Call this function to register func as a builtin function.

This function can be called at module-level scope to register func as a builtin function. A builtin function will be converted to a CallFunction Expr in tracing:

def my_func(x, y):
    return x + y

import megengine.traced_module as tm
tm.wrap(my_func)

This function can also equivalently be used as a decorator:

@tm.wrap
def my_func(x, y):
    return x + y
Parameters:

func (Callable) – the function of the global function to insert into the graph when it’s called.