def _fn_to_typed_def(fn, method=False): schema = annotations.get_signature(fn) ast = get_jit_ast(fn) if schema: typed_def = torch._C._pack_typed_def(ast, schema[0], schema[1], method) else: typed_def = torch._C.TypedDef(ast) return typed_def
def script(fn, optimize=True, _frames_up=0): rcb = createResolutionCallback(_frames_up + 1) ast = get_jit_ast(fn, is_method=False) graph = _jit_script_compile(ast, rcb) mod = ScriptModule() mod._create_method_from_graph('forward', graph) # TODO: refactor everything so we're not 1) creating a ScriptModule # 2) Throwing everything away except for the graph 3) Creating a new # ScriptModule and dumping that graph in 4) Re-populating the schema # because it was lost doing the previous mod.__getattr__('forward').forward_schema(ast, False) # Forward docstrings mod.__doc__ = fn.__doc__ return mod
def script_method(fn): # NOTE: we need to traverse two frames here because the meta-class frame # for ScriptModule will be present, as opposed to invoking @script on a # a function or invoking define() on a CompilationUnit. # The stack will look like: # # 0. createResolutionCallback() # 1. script_method() # 2. ScriptModule metaclass frame # 3. Surrounding scope # # createResolutionCallback internally adds 1 to get us to the scope of this # function (the calling function). Adding 2 gets us to the proper surrounding scope. return ScriptMethodStub(createResolutionCallback(frames_up=2), get_jit_ast(fn), fn)
def script_method(fn): return ScriptMethodStub(createResolutionCallback(frames_up=1), get_jit_ast(fn), fn)
def _script_graph(fn, _frames_up=0): rcb = createResolutionCallback(_frames_up + 1) ast = get_jit_ast(fn) return _jit_script_compile(ast, rcb)
def script_method(fn): return ScriptMethodStub(createResolutionCallback(), get_jit_ast(fn))
def _script_graph(fn, frame_id=2): rcb = createResolutionCallback(frame_id) ast = get_jit_ast(fn) return _jit_script_compile(ast, rcb)
def script(fn): rcb = createResolutionCallback() ast = get_jit_ast(fn) graph = _jit_script_compile(ast, rcb) return torch._C.GraphExecutor(graph, True)
def script_method(fn): return ScriptMethodStub(createResolutionCallback(frames_up=1), get_jit_ast(fn), fn)
def _script_graph(fn, _frames_up=0): rcb = createResolutionCallback(_frames_up + 1) ast = get_jit_ast(fn) return _jit_script_compile(ast, rcb)