neulab/docprompting-conala huggingface dataset:
set_result(result) [source]
 
Set the result for this Future, which will mark this Future as completed and trigger all attached callbacks. Note that a Future cannot be marked completed twice.  Parameters 
result (object) – the result object of this Future.    Example::

>>> import threading
>>> import time
>>> import torch
>>>
>>> def slow_set_future(fut, value):
>>>     time.sleep(0.5)
>>>     fut.set_result(value)
>>>
>>> fut = torch.futures.Future()
>>> t = threading.Thread(
>>>     target=slow_set_future,
>>>     args=(fut, torch.ones(2) * 3)
>>> )
>>> t.start()
>>>
>>> print(fut.wait())  # tensor([3., 3.])
>>> t.join()
then(callback) [source]
 
Append the given callback function to this Future, which will be run when the Future is completed. Multiple callbacks can be added to the same Future, and will be invoked in the same order as they were added. The callback must take one argument, which is the reference to this Future. The callback function can use the Future.wait() API to get the value. Note that if this Future is already completed, the given callback will be run immediately inline.  Parameters 
callback (Callable) – a Callable that takes this Future as the only argument.  Returns 
A new Future object that holds the return value of the callback and will be marked as completed when the given callback finishes.    Example::

>>> import torch
>>>
>>> def callback(fut):
>>>     print(f"RPC return value is {fut.wait()}.")
>>>
>>> fut = torch.futures.Future()
>>> # The inserted callback will print the return value when
>>> # receiving the response from "worker1"
>>> cb_fut = fut.then(callback)
>>> chain_cb_fut = cb_fut.then(
>>>     lambda x : print(f"Chained cb done. {x.wait()}")
>>> )
>>> fut.set_result(5)
>>>
>>> # Outputs are:
>>> # RPC return value is 5.
>>> # Chained cb done. None
wait() [source]
 
Block until the value of this Future is ready.  Returns 
The value held by this Future. If the function (callback or RPC) creating the value has thrown an error, this wait method will also throw an error.
torch.futures.wait_all(futures) [source]
 
Waits for all provided futures to be complete, and returns the list of completed values.  Parameters 
futures (list) – a list of Future object.  Returns 
A list of the completed Future results. This method will throw an error if wait on any Future throws.
call_function(the_function, args=None, kwargs=None, type_expr=None) [source]
 
Insert a call_function Node into the Graph. A call_function node represents a call to a Python callable, specified by the_function. the_function can be  Parameters 
 
the_function (Callable[.., Any]) – The function to be called. Can be any PyTorch operator, Python function, or member of the builtins or operator namespaces. 
args (Optional[Tuple[Argument, ..]]) – The positional arguments to be passed to the called function. 
kwargs (Optional[Dict[str, Argument]]) – The keyword arguments to be passed to the called function 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.    Returns The newly created and inserted call_function node.  Note The same insertion point and type expression rules apply for this method as Graph.create_node().
call_method(method_name, args=None, kwargs=None, type_expr=None) [source]
 
Insert a call_method Node into the Graph. A call_method node represents a call to a given method on the 0th element of args.  Parameters 
 
method_name (str) – The name of the method to apply to the self argument. For example, if args[0] is a Node representing a Tensor, then to call relu() on that Tensor, pass relu to method_name. 
args (Optional[Tuple[Argument, ..]]) – The positional arguments to be passed to the called method. Note that this should include a self argument. 
kwargs (Optional[Dict[str, Argument]]) – The keyword arguments to be passed to the called method 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly created and inserted call_method node.    Note The same insertion point and type expression rules apply for this method as Graph.create_node().
call_module(module_name, args=None, kwargs=None, type_expr=None) [source]
 
Insert a call_module Node into the Graph. A call_module node represents a call to the forward() function of a Module in the Module hierarchy.  Parameters 
 
module_name (str) – The qualified name of the Module in the Module hierarchy to be called. For example, if the traced Module has a submodule named foo, which has a submodule named bar, the qualified name foo.bar should be passed as module_name to call that module. 
args (Optional[Tuple[Argument, ..]]) – The positional arguments to be passed to the called method. Note that this should not include a self argument. 
kwargs (Optional[Dict[str, Argument]]) – The keyword arguments to be passed to the called method 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly-created and inserted call_module node.    Note The same insertion point and type expression rules apply for this method as Graph.create_node().
create_node(op, target, args=None, kwargs=None, name=None, type_expr=None) [source]
 
Create a Node and add it to the Graph at the current insert-point. Note that the current insert-point can be set via Graph.inserting_before() and Graph.inserting_after().  Parameters 
 
op (str) – the opcode for this Node. One of ‘call_function’, ‘call_method’, ‘get_attr’, ‘call_module’, ‘placeholder’, or ‘output’. The semantics of these opcodes are described in the Graph docstring. 
args (Optional[Tuple[Argument, ..]]) – is a tuple of arguments to this node. 
kwargs (Optional[Dict[str, Argument]]) – the kwargs of this Node 
name (Optional[str]) – an optional string name for the Node. This will influence the name of the value assigned to in the Python generated code. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly-created and inserted node.
erase_node(to_erase) [source]
 
Erases a Node from the Graph. Throws an exception if there are still users of that node in the Graph.  Parameters 
to_erase (Node) – The Node to erase from the Graph.
get_attr(qualified_name, type_expr=None) [source]
 
Insert a get_attr node into the Graph. A get_attr Node represents the fetch of an attribute from the Module hierarchy.  Parameters 
 
qualified_name (str) – the fully-qualified name of the attribute to be retrieved. For example, if the traced Module has a submodule named foo, which has a submodule named bar, which has an attribute named baz, the qualified name foo.bar.baz should be passed as qualified_name. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly-created and inserted get_attr node.    Note The same insertion point and type expression rules apply for this method as Graph.create_node.
graph_copy(g, val_map) [source]
 
Copy all nodes from a given graph into self.  Parameters 
 
g (Graph) – The source graph from which to copy Nodes. 
val_map (Dict[Node, Node]) – a dictionary that will be populated with a mapping from nodes in g to nodes in self. Note that val_map can be passed in with values in it already to override copying of certain values.   Returns 
The value in self that is now equivalent to the output value in g, if g had an output node. None otherwise.
inserting_after(n=None) [source]
 
Set the point at which create_node and companion methods will insert into the graph. When used within a ‘with’ statement, this will temporary set the insert point and then restore it when the with statement exits: with g.inserting_after(n):
    ... # inserting after node n
... # insert point restored to what it was previously
g.inserting_after(n) #  set the insert point permanently
  Parameters 
n (Optional[Node]) – The node before which to insert. If None this will insert after the beginning of the entire graph.  Returns 
A resource manager that will restore the insert point on __exit__.
inserting_before(n=None) [source]
 
Set the point at which create_node and companion methods will insert into the graph. When used within a ‘with’ statement, this will temporary set the insert point and then restore it when the with statement exits: with g.inserting_before(n):
    ... # inserting before node n
... # insert point restored to what it was previously
g.inserting_before(n) #  set the insert point permanently
  Parameters 
n (Optional[Node]) – The node before which to insert. If None this will insert before the beginning of the entire graph.  Returns 
A resource manager that will restore the insert point on __exit__.
lint(root=None) [source]
 
Runs various checks on this Graph to make sure it is well-formed. In particular: - Checks Nodes have correct ownership (owned by this graph) - Checks Nodes appear in topological order - If root is provided, checks that targets exist in root  Parameters 
root (Optional[torch.nn.Module]) – The root module with which to check for targets. This is equivalent to the root argument that is passed when constructing a GraphModule.
property nodes  
Get the list of Nodes that constitute this Graph. Note that this Node list representation is a doubly-linked list. Mutations during iteration (e.g. delete a Node, add a Node) are safe.  Returns 
A doubly-linked list of Nodes. Note that reversed can be called on this list to switch iteration order.
node_copy(node, arg_transform=<function Graph.<lambda>>) [source]
 
Copy a node from one graph into another. arg_transform needs to transform arguments from the graph of node to the graph of self. Example: # Copying all the nodes in `g` into `new_graph`
g : torch.fx.Graph = ...
new_graph = torch.fx.graph()
value_remap = {}
for node in g.nodes:
    value_remap[node] = new_graph.node_copy(node, lambda n : value_remap[n])
  Parameters 
 
node (Node) – The node to copy into self. 
arg_transform (Callable[[Node], Argument]) – A function that transforms Node arguments in node’s args and kwargs into the equivalent argument in self. In the simplest case, this should retrieve a value out of a table mapping Nodes in the original graph to self.
output(result, type_expr=None) [source]
 
Insert an output Node into the Graph. An output node represents a return statement in Python code. result is the value that should be returned.  Parameters 
 
result (Argument) – The value to be returned. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.     Note The same insertion point and type expression rules apply for this method as Graph.create_node.
placeholder(name, type_expr=None) [source]
 
Insert a placeholder node into the Graph. A placeholder represents a function input.  Parameters 
 
name (str) – A name for the input value. This corresponds to the name of the positional argument to the function this Graph represents. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have. This is needed in some cases for proper code generation (e.g. when the function is used subsequently in TorchScript compilation).     Note The same insertion point and type expression rules apply for this method as Graph.create_node.
print_tabular() [source]
 
Prints the intermediate representation of the graph in tabular format.
python_code(root_module) [source]
 
Turn this Graph into valid Python code.  Parameters 
root_module (str) – The name of the root module on which to look-up qualified name targets. This is usually ‘self’.  Returns 
The string source code generated from this Graph.
class torch.fx.GraphModule(root, graph, class_name='GraphModule') [source]
 
GraphModule is an nn.Module generated from an fx.Graph. Graphmodule has a graph attribute, as well as code and forward attributes generated from that graph.  Warning When graph is reassigned, code and forward will be automatically regenerated. However, if you edit the contents of the graph without reassigning the graph attribute itself, you must call recompile() to update the generated code.   
__init__(root, graph, class_name='GraphModule') [source]
 
Construct a GraphModule.  Parameters 
 
root (Union[torch.nn.Module, Dict[str, Any]) – root can either be an nn.Module instance or a Dict mapping strings to any attribute type. In the case that root is a Module, any references to Module-based objects (via qualified name) in the Graph’s Nodes’ target field will be copied over from the respective place within root’s Module hierarchy into the GraphModule’s module hierarchy. In the case that root is a dict, the qualified name found in a Node’s target will be looked up directly in the dict’s keys. The object mapped to by the Dict will be copied over into the appropriate place within the GraphModule’s module hierarchy. 
graph (Graph) – graph contains the nodes this GraphModule should use for code generation 
name (str) – name denotes the name of this GraphModule for debugging purposes. If it’s unset, all error messages will report as originating from GraphModule. It may be helpful to set this to root’s original name or a name that makes sense within the context of your transform.    
  
property code  
Return the Python code generated from the Graph underlying this GraphModule. 
  
property graph  
Return the Graph underlying this GraphModule 
  
recompile() [source]
 
Recompile this GraphModule from its graph attribute. This should be called after editing the contained graph, otherwise the generated code of this GraphModule will be out of date. 
  
to_folder(folder, module_name='FxModule') [source]
 
Dumps out module to folder with module_name so that it can be imported with from <folder> import <module_name>  Parameters 
 
folder (Union[str, os.PathLike]) – The folder to write the code out to 
module_name (str) – Top-level name to use for the Module while writing out the code
recompile() [source]
 
Recompile this GraphModule from its graph attribute. This should be called after editing the contained graph, otherwise the generated code of this GraphModule will be out of date.
to_folder(folder, module_name='FxModule') [source]
 
Dumps out module to folder with module_name so that it can be imported with from <folder> import <module_name>  Parameters 
 
folder (Union[str, os.PathLike]) – The folder to write the code out to 
module_name (str) – Top-level name to use for the Module while writing out the code
__init__(root, graph, class_name='GraphModule') [source]
 
Construct a GraphModule.  Parameters 
 
root (Union[torch.nn.Module, Dict[str, Any]) – root can either be an nn.Module instance or a Dict mapping strings to any attribute type. In the case that root is a Module, any references to Module-based objects (via qualified name) in the Graph’s Nodes’ target field will be copied over from the respective place within root’s Module hierarchy into the GraphModule’s module hierarchy. In the case that root is a dict, the qualified name found in a Node’s target will be looked up directly in the dict’s keys. The object mapped to by the Dict will be copied over into the appropriate place within the GraphModule’s module hierarchy. 
graph (Graph) – graph contains the nodes this GraphModule should use for code generation 
name (str) – name denotes the name of this GraphModule for debugging purposes. If it’s unset, all error messages will report as originating from GraphModule. It may be helpful to set this to root’s original name or a name that makes sense within the context of your transform.
call_function(target, args, kwargs) [source]
 
Execute a call_function node and return the result.  Parameters 
 
target (Target) – The call target for this node. See Node for details on semantics 
args (Tuple) – Tuple of positional args for this invocation 
kwargs (Dict) – Dict of keyword arguments for this invocation     Return

Any: The value returned by the function invocation
call_method(target, args, kwargs) [source]
 
Execute a call_method node and return the result.  Parameters 
 
target (Target) – The call target for this node. See Node for details on semantics 
args (Tuple) – Tuple of positional args for this invocation 
kwargs (Dict) – Dict of keyword arguments for this invocation     Return

Any: The value returned by the method invocation
call_module(target, args, kwargs) [source]
 
Execute a call_module node and return the result.  Parameters 
 
target (Target) – The call target for this node. See Node for details on semantics 
args (Tuple) – Tuple of positional args for this invocation 
kwargs (Dict) – Dict of keyword arguments for this invocation     Return

Any: The value returned by the module invocation
fetch_args_kwargs_from_env(n) [source]
 
Fetch the concrete values of args and kwargs of node n from the current execution environment.  Parameters 
n (Node) – The node for which args and kwargs should be fetched.  Returns 
args and kwargs with concrete values for n.  Return type 
Tuple[Tuple, Dict]
fetch_attr(target) [source]
 
Fetch an attribute from the Module hierarchy of self.module.  Parameters 
target (str) – The fully-qualfiied name of the attribute to fetch  Returns 
The value of the attribute.  Return type 
Any
get_attr(target, args, kwargs) [source]
 
Execute a get_attr node. Will retrieve an attribute value from the Module hierarchy of self.module.  Parameters 
 
target (Target) – The call target for this node. See Node for details on semantics 
args (Tuple) – Tuple of positional args for this invocation 
kwargs (Dict) – Dict of keyword arguments for this invocation   Returns 
The value of the attribute that was retrieved  Return type 
Any
map_nodes_to_values(args, n) [source]
 
Recursively descend through args and look up the concrete value for each Node in the current execution environment.  Parameters 
 
args (Argument) – Data structure within which to look up concrete values 
n (Node) – Node to which args belongs. This is only used for error reporting.
output(target, args, kwargs) [source]
 
Execute an output node. This really just retrieves the value referenced by the output node and returns it.  Parameters 
 
target (Target) – The call target for this node. See Node for details on semantics 
args (Tuple) – Tuple of positional args for this invocation 
kwargs (Dict) – Dict of keyword arguments for this invocation   Returns 
The return value referenced by the output node  Return type 
Any
placeholder(target, args, kwargs) [source]
 
Execute a placeholder node. Note that this is stateful: Interpreter maintains an internal iterator over arguments passed to run and this method returns next() on that iterator.  Parameters 
 
target (Target) – The call target for this node. See Node for details on semantics 
args (Tuple) – Tuple of positional args for this invocation 
kwargs (Dict) – Dict of keyword arguments for this invocation   Returns 
The argument value that was retrieved.  Return type 
Any
run(*args, initial_env=None) [source]
 
Run module via interpretation and return the result.  Parameters 
 
*args – The arguments to the Module to run, in positional order 
initial_env (Optional[Dict[Node, Any]]) – An optional starting environment for execution. This is a dict mapping Node to any value. This can be used, for example, to pre-populate results for certain Nodes so as to do only partial evaluation within the interpreter.   Returns 
The value returned from executing the Module  Return type 
Any
run_node(n) [source]
 
Run a specific node n and return the result. Calls into placeholder, get_attr, call_function, call_method, call_module, or output depending on node.op  Parameters 
n (Node) – The Node to execute  Returns 
The result of executing n  Return type 
Any
property all_input_nodes  
Return all Nodes that are inputs to this Node. This is equivalent to iterating over args and kwargs and only collecting the values that are Nodes.  Returns 
List of Nodes that appear in the args and kwargs of this Node, in that order.
append(x) [source]
 
Insert x after this node in the list of nodes in the graph. Equvalent to self.next.prepend(x)  Parameters 
x (Node) – The node to put after this node. Must be a member of the same graph.
property args  
The tuple of arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information. Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.
property kwargs  
The dict of keyword arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information. Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.
property next  
Returns the next Node in the linked list of Nodes.  Returns 
The next Node in the linked list of Nodes.
prepend(x) [source]
 
Insert x before this node in the list of nodes in the graph. Example: Before: p -> self
        bx -> x -> ax
After:  p -> x -> self
        bx -> ax
  Parameters 
x (Node) – The node to put before this node. Must be a member of the same graph.
property prev  
Returns the previous Node in the linked list of Nodes.  Returns 
The previous Node in the linked list of Nodes.
replace_all_uses_with(replace_with) [source]
 
Replace all uses of self in the Graph with the Node replace_with.  Parameters 
replace_with (Node) – The node to replace all uses of self with.  Returns 
The list of Nodes on which this change was made.
class torch.fx.Proxy(node, tracer=None) [source]
 
Proxy objects are Node wrappers that flow through the program during symbolic tracing and record all the operations (torch function calls, method calls, operators) that they touch into the growing FX Graph. If you’re doing graph transforms, you can wrap your own Proxy method around a raw Node so that you can use the overloaded operators to add additional things to a Graph.
torch.fx.symbolic_trace(root, concrete_args=None) [source]
 
Symbolic tracing API Given an nn.Module or function instance root, this function will return a GraphModule constructed by recording operations seen while tracing through root.  Parameters 
 
root (Union[torch.nn.Module, Callable]) – Module or function to be traced and converted into a Graph representation. 
concrete_args (Optional[Dict[str, any]]) – Concrete arguments that should not be treated as Proxies.   Returns 
a Module created from the recorded operations from root.  Return type 
GraphModule
call_module(m, forward, args, kwargs) [source]
 
Method that specifies the behavior of this Tracer when it encounters a call to an nn.Module instance. By default, the behavior is to check if the called module is a leaf module via is_leaf_module. If it is, emit a call_module node referring to m in the Graph. Otherwise, call the Module normally, tracing through the operations in its forward function. This method can be overridden to–for example–create nested traced GraphModules, or any other behavior you would want while tracing across Module boundaries. Module boundaries.  Parameters 
 
m (Module) – The module for which a call is being emitted 
forward (Callable) – The forward() method of the Module to be invoked 
args (Tuple) – args of the module callsite 
kwargs (Dict) – kwargs of the module callsite   Returns 
The return value from the Module call. In the case that a call_module node was emitted, this is a Proxy value. Otherwise, it is whatever value was returned from the Module invocation.
create_arg(a) [source]
 
A method to specify the behavior of tracing when preparing values to be used as arguments to nodes in the Graph. By default, the behavior includes:  Iterate through collection types (e.g. tuple, list, dict) and recursively call create_args on the elements. Given a Proxy object, return a reference to the underlying IR Node
 
Given a non-Proxy Tensor object, emit IR for various cases:  For a Parameter, emit a get_attr node referring to that Parameter For a non-Parameter Tensor, store the Tensor away in a special attribute referring to that attribute.    This method can be overridden to support more types.  Parameters 
a (Any) – The value to be emitted as an Argument in the Graph.  Returns 
The value a converted into the appropriate Argument