Add docstring for Proxy (#50145)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/50145

Test Plan: Imported from OSS

Reviewed By: pbelevich

Differential Revision: D25854281

Pulled By: ansley

fbshipit-source-id: d7af6fd6747728ef04e86fbcdeb87cb0508e1fd8
This commit is contained in:
Ansley Ussery 2021-01-11 13:46:10 -08:00 committed by Facebook GitHub Bot
parent 3d263d1928
commit 080a097935
2 changed files with 13 additions and 7 deletions

View File

@ -28,3 +28,5 @@ API Reference
.. autoclass:: torch.fx.Tracer
:members:
.. autoclass:: torch.fx.Proxy

View File

@ -112,17 +112,21 @@ class GraphAppendingTracer(TracerBase):
class TraceError(ValueError):
pass
# Proxy objects are stand-in values for normal values in a PyTorch computation.
# Instead of performing compute they record computation into Graph.
# Each proxy wraps the Node instance that represents the expression that define the
# value.
class Proxy:
"""
``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``.
"""
def __init__(self, node: Node, tracer: 'Optional[TracerBase]' = None):
if tracer is None:
# this allows you to create a proxy object around a raw node
# so that if you are doing graph transforms you can use the overloaded operators
# to add additional things to a graph.
# This allows you to create a Proxy object around a raw Node
tracer = GraphAppendingTracer(node.graph)
self.tracer = tracer
self.node = node