def script(pyfunc): """Decorate a python function function as hybrid script. The hybrid function support emulation mode and parsing to the internal language IR. Returns ------- hybrid_func : function A decorated hybrid script function. """ # pylint: disable=import-outside-toplevel, missing-docstring def wrapped_func(func, *args, **kwargs): from .utils import _is_tvm_arg_types if _is_tvm_arg_types(args): src = _pruned_source(func) closure_vars = inspect.getclosurevars(func).nonlocals closure_vars.update(inspect.getclosurevars(func).globals) return source_to_op(src, args, func.__globals__, closure_vars) from .runtime import _enter_hybrid_runtime, _restore_runtime intersect = _enter_hybrid_runtime(func) value = func(*args, **kwargs) _restore_runtime(func, intersect) return value return decorate(pyfunc, wrapped_func)
def __call__(self, fdecl): def tagged_fdecl(func, *args, **kwargs): with self: return func(*args, **kwargs) return decorate(fdecl, tagged_fdecl)