MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / wrap

Function wrap

imperative/python/megengine/traced_module/traced_module.py:2381–2414  ·  view source on GitHub ↗

r"""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 :class:`CallFunction` Expr in tracing:: def my_func(x, y): return x

(func: Callable)

Source from the content-addressed store, hash-verified

2379
2380
2381def wrap(func: Callable):
2382 r"""Call this function to register ``func`` as a builtin function.
2383
2384 This function can be called at module-level scope to register ``func`` as a builtin function.
2385 A builtin function will be converted to a :class:`CallFunction` Expr in tracing::
2386
2387 def my_func(x, y):
2388 return x + y
2389
2390 import megengine.traced_module as tm
2391 tm.wrap(my_func)
2392
2393 This function can also equivalently be used as a decorator::
2394
2395 @tm.wrap
2396 def my_func(x, y):
2397 return x + y
2398
2399 Args:
2400 func: the function of the global function to insert into the graph when it's called.
2401 """
2402 USER_REGISTERED_FUNCTION.append((func.__module__, func.__qualname__))
2403 assert callable(func), "func must be a callable"
2404 assert hasattr(func, "__code__")
2405 fn_name = func.__code__.co_name
2406 currentframe = inspect.currentframe()
2407 assert currentframe is not None
2408 f = currentframe.f_back
2409 assert f is not None
2410 assert (
2411 f.f_code.co_name == "<module>"
2412 ), "wrap must be called at the top level of a module"
2413 Patcher._builtin_functions.append((f.f_globals, fn_name))
2414 return func
2415
2416
2417def _register_all_builtin_module():

Callers 2

__call__Method · 0.85
make_backward_graphFunction · 0.85

Calls 1

appendMethod · 0.45

Tested by

no test coverage detected