/usr/local/lib64/python3.6/site-packages/torch/fx/__pycache__
Edit: /usr/local/lib64/python3.6/site-packages/torch/fx/__pycache__/interpreter.cpython-36.pyc (17660B)
3
EgmH @ s d dl mZ d dlmZ d dlmZmZmZmZm Z d dl
m Z d dl m
Z
d dlmZ ddlmZmZmZmZmZmZmZ ed d
G d d d Zed d
G d
d deZdS ) ) GraphModule)Graph)ArgumentNodeTargetmap_arg
map_aggregate)Proxy)Tracer)
compatibility )AnyDictIteratorListOptionalTupleUnionT)is_backward_compatiblec @ s e Zd ZdZeddd&eedddZeddddee e
e f e d d
d Z edde
e d d
dZ
edddeedf e ee f e dddZedddeedf e ee f e dddZedddeedf e ee f e dddZedddeedf e ee f e dddZedddeedf e ee f e dddZedddeedf e ee f e dd d Zedded d d Zedde
eee f d d!d"Zeddee
ed#d$d%ZdS )' Interpretera
An Interpreter executes an FX graph Node-by-Node. This pattern
can be useful for many things, including writing code
transformations as well as analysis passes.
Methods in the Interpreter class can be overridden to customize
the behavior of execution. The map of overrideable methods
in terms of call hierarchy::
run()
+-- run_node
+-- placeholder()
+-- get_attr()
+-- call_function()
+-- call_method()
+-- call_module()
+-- output()
Example:
Suppose we want to swap all instances of ``torch.neg`` with
``torch.sigmoid`` and vice versa (including their ``Tensor``
method equivalents). We could subclass Interpreter like so::
class NegSigmSwapInterpreter(Interpreter):
def call_function(self, target : Target,
args : Tuple, kwargs : Dict) -> Any:
if target == torch.sigmoid:
return torch.neg(*args, **kwargs)
return super().call_function(n)
def call_method(self, target : Target,
args : Tuple, kwargs : Dict) -> Any:
if target == 'neg':
call_self, *args_tail = args
return call_self.sigmoid(*args_tail, **kwargs)
return super().call_method(n)
def fn(x):
return torch.sigmoid(x).neg()
gm = torch.fx.symbolic_trace(fn)
input = torch.randn(3, 4)
result = NegSigmSwapInterpreter(gm).run(input)
torch.testing.assert_allclose(result, torch.neg(input).sigmoid())
Args:
module (GraphModule): The module to be executed
garbage_collect_values (bool): Whether to delete values after their last
use within the Module's execution. This ensures optimal memory usage during
execution. This can be disabled to, for example, examine all of the intermediate
values in the execution by looking at the ``Interpreter.env`` attribute.
T)r )modulegarbage_collect_valuesc s t |tst|_tjj _i _|_jri i _ t
t
dfdd xBt jj j
D ]0 t j fdd t j fdd qfW d S )N)nuserc s( | kr$| | < j j|g j| d S )N)user_to_last_uses
setdefaultappend)r r )node_to_last_useself @/usr/local/lib64/python3.6/site-packages/torch/fx/interpreter.pyregister_last_usesQ s z0Interpreter.__init__.
.register_last_usesc s
| S )Nr )r )noder! r r W s z&Interpreter.__init__..c s
| S )Nr )r )r" r! r r r# X s )
isinstancer AssertionErrorr dictZ
named_modulesZ
submodulesenvr r r reversedgraphnodesr argskwargs)r r r r )r" r r! r r __init__@ s zInterpreter.__init__N) initial_env)r. returnc G s |r|ni | _ t|| _xl| jjjD ]^}|| j kr4q$| j|| j |< | jrjx | jj |g D ] }| j |= qZW |j
dkr$| j | }|S q$W dS )a:
Run `module` via interpretation and return the result.
Args:
*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:
Any: The value returned from executing the Module
outputN) r' iter args_iterr r) r* run_noder r getop)r r. r+ r" Z to_deleteZ
output_valr r r runZ s
zInterpreter.run)r r/ c C s@ | j |\}}t|ts tt|ts*tt| |j|j||S )aB
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``
Args:
n (Node): The Node to execute
Returns:
Any: The result of executing ``n``
)fetch_args_kwargs_from_envr$ tupler% r& getattrr5 target)r r r+ r, r r r r3 s zInterpreter.run_noder .)r: r+ r, r/ c C s0 t |tst|jdr"t| jS t| jS dS )a
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.
Args:
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:
Any: The argument value that was retrieved.
*N)r$ strr%
startswithlistr2 next)r r: r+ r, r r r placeholder s
zInterpreter.placeholderc C s t |tst| j|S )a2
Execute a ``get_attr`` node. Will retrieve an attribute
value from the ``Module`` hierarchy of ``self.module``.
Args:
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 of the attribute that was retrieved
)r$ r< r%
fetch_attr)r r: r+ r, r r r get_attr s zInterpreter.get_attrc C s t |t st|||S )a
Execute a ``call_function`` node and return the result.
Args:
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
)r$ r< r% )r r: r+ r, r r r
call_function s zInterpreter.call_functionc C s& |^}}t |tstt||||S )a
Execute a ``call_method`` node and return the result.
Args:
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
)r$ r< r% r9 )r r: r+ r, Zself_objZ args_tailr r r call_method s zInterpreter.call_methodc C s" t |tst| j|}|||S )a
Execute a ``call_module`` node and return the result.
Args:
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
)r$ r< r% rA )r r: r+ r, submodr r r call_module s
zInterpreter.call_modulec C s |d S )a6
Execute an ``output`` node. This really just retrieves
the value referenced by the ``output`` node and returns it.
Args:
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 return value referenced by the output node
r r )r r: r+ r, r r r r0 s zInterpreter.output)r: c C sZ |j d}| j}xDt|D ]8\}}t||sHtddj|d| t||}qW |S )z
Fetch an attribute from the ``Module`` hierarchy of ``self.module``.
Args:
target (str): The fully-qualfiied name of the attribute to fetch
Return:
Any: The value of the attribute.
.z#Node referenced nonexistent target N)splitr enumeratehasattr RuntimeErrorjoinr9 )r r: Z target_atomsZattr_itriZatomr r r rA s
zInterpreter.fetch_attrc C s@ | j |j|}t|ts t| j |j|}t|ts8t||fS )aP
Fetch the concrete values of ``args`` and ``kwargs`` of node ``n``
from the current execution environment.
Args:
n (Node): The node for which ``args`` and ``kwargs`` should be fetched.
Return:
Tuple[Tuple, Dict]: ``args`` and ``kwargs`` with concrete values for ``n``.
)map_nodes_to_valuesr+ r$ r8 r% r, r& )r r r+ r, r r r r7 * s
z&Interpreter.fetch_args_kwargs_from_env)r+ r r/ c s t td fdd }t||S )aV
Recursively descend through ``args`` and look up the concrete value
for each ``Node`` in the current execution environment.
Args:
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.
)n_argr/ c s* | j kr td d| dj | S )NzNode z referenced nonexistent value z*! Run Graph.lint() to diagnose such issues)r' rK )rO )r r r r load_argG s
z1Interpreter.map_nodes_to_values..load_arg)r r
r )r r+ r rP r )r r r rN < s z Interpreter.map_nodes_to_values)T)__name__
__module__ __qualname____doc__r r boolr- r r r r
r6 r3 r r r< r@ rB rC rD rF r0 rA r7 rN r r r r r s2 6$'&&&&&&r c s e Zd ZdZedd fddZedddeedf ee e
f e dd d
Z edddeedf ee e
f e dd d Z
edddeedf ee e
f e
dd
dZedddeedf ee e
f e
dddZedded fdd Z ZS ) Transformera9
``Transformer`` is a special type of interpreter that produces a
new ``Module``. It exposes a ``transform()`` method that returns
the transformed ``Module``. ``Transformer`` does not require
arguments to run, as ``Interpreter`` does. ``Transformer`` works
entirely symbolically.
Example:
Suppose we want to swap all instances of ``torch.neg`` with
``torch.sigmoid`` and vice versa (including their ``Tensor``
method equivalents). We could subclass ``Transformer`` like so::
class NegSigmSwapXformer(Transformer):
def call_function(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
if target == torch.sigmoid:
return torch.neg(*args, **kwargs)
return super().call_function(n)
def call_method(self, target : 'Target', args : Tuple[Argument, ...], kwargs : Dict[str, Any]) -> Any:
if target == 'neg':
call_self, *args_tail = args
return call_self.sigmoid(*args_tail, **kwargs)
return super().call_method(n)
def fn(x):
return torch.sigmoid(x).neg()
gm = torch.fx.symbolic_trace(fn)
transformed : torch.nn.Module = NegSigmSwapXformer(gm).transform()
input = torch.randn(3, 4)
torch.testing.assert_allclose(transformed(input), torch.neg(input).sigmoid())
Args:
module (GraphModule): The ``Module`` to be transformed.
T)r c s< t j| t | _G dd dt}|| j| _|| j_d S )Nc s0 e Zd Zed fdd ZedddZ ZS )z/Transformer.__init__..TransformerTracer)r) c s t j || _d S )N)superr- r) )r r) ) __class__r r r- | s
z8Transformer.__init__..TransformerTracer.__init__)r/ c S s dS )NTr )r ___r r r is_leaf_module s z>Transformer.__init__..TransformerTracer.is_leaf_module)rQ rR rS r r- rU r[
__classcell__r r )rX r TransformerTracer{ s r] )rW r- r new_graphr
tracerroot)r r r] )rX r r r- v s
zTransformer.__init__r .)r: r+ r, r/ c C s" t |tstt| jj|| jS )a
Execute a ``placeholder`` node. In ``Transformer``, this is
overridden to insert a new ``placeholder`` into the output
graph.
Args:
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
)r$ r< r% r r^ r@ r_ )r r: r+ r, r r r r@ s zTransformer.placeholderc C s" t |tstt| jj|| jS )a
Execute a ``get_attr`` node. In ``Transformer``, this is
overridden to insert a new ``get_attr`` node into the output
graph.
Args:
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
)r$ r< r% r r^ rB r_ )r r: r+ r, r r r rB s zTransformer.get_attrc C s, t |tst| j|}| jj||j||S )N)r$ r< r% rA r_ rF Zforward)r r: r+ r, rE r r r rF s
zTransformer.call_modulec C s | j jd|||S )NrC )r_ Z create_proxy)r r: r+ r, r r r rC s zTransformer.call_function)r/ c sJ t j }|dk r.strip_proxy) rW r6 r r r r
r^ r0 r r r )r resultrb )rX r r transform s
zTransformer.transform)rQ rR rS rT r r- r r r r< r
r r@ rB rF rC r rd r\ r r )rX r rV N s &&&&&rV N)Z graph_moduler r) r r" r r r r r proxyr Z_symbolic_tracer
Z_compatibilityr typingr
r r r r r r r rV r r r r s $ G