('batch', 'channels', 'H', 'W')
     >>> renamed_imgs = imgs.rename(None)
     >>> renamed_imgs.names
     (None, None, None, None)

     >>> renamed_imgs = imgs.rename('batch', 'channel', 'height', 'width')
     >>> renamed_imgs.names
     ('batch', 'channel', 'height', 'width')

  Warning:

    The named tensor API is experimental and subject to change.

rename_(names, *rename_map)
  In-place version of "rename()".

refine_names(*names)
  Refines the dimension names of "self" according to "names".

  Refining is a special case of renaming that "lifts" unnamed
  dimensions. A "None" dim can be refined to have any name; a
  named dim can only be refined to have the same name.

  Because named tensors can coexist with unnamed tensors, refining
  names gives a nice way to write named-tensor-aware code that
  works with both named and unnamed tensors.

  "names" may contain up to one Ellipsis ("..."). The Ellipsis is


expanded greedily; it is expanded in-place to fill "names" to
      the same length as "self.dim()" using names from the
      corresponding indices of "self.names".
  Python 2 does not support Ellipsis but one may use a string
  literal instead ("'...'").

  Parameters:
     **names** (*iterable of str*) -- The desired names of the
     output tensor. May contain up to one Ellipsis.

  Examples:

     >>> imgs = torch.randn(32, 3, 128, 128)
     >>> named_imgs = imgs.refine_names('N', 'C', 'H', 'W')
     >>> named_imgs.names
     ('N', 'C', 'H', 'W')

     >>> tensor = torch.randn(2, 3, 5, 7, 11)
     >>> tensor = tensor.refine_names('A', ..., 'B', 'C')
     >>> tensor.names
     ('A', None, None, 'B', 'C')

  Warning:

    The named tensor API is experimental and subject to change.

align_as(other) -> Tensor
  Permutes the dimensions of the "self" tensor to match the


dimension order in the "other" tensor, adding size-one dims for
      any new names.
  This operation is useful for explicit broadcasting by names (see
  examples).

  All of the dims of "self" must be named in order to use this
  method. The resulting tensor is a view on the original tensor.

  All dimension names of "self" must be present in "other.names".
  "other" may contain named dimensions that are not in
  "self.names"; the output tensor has a size-one dimension for
  each of those new names.

  To align a tensor to a specific order, use "align_to()".

  Examples:

     # Example 1: Applying a mask
     >>> mask = torch.randint(2, [127, 128], dtype=torch.bool).refine_names('W', 'H')
     >>> imgs = torch.randn(32, 128, 127, 3, names=('N', 'H', 'W', 'C'))
     >>> imgs.masked_fill_(mask.align_as(imgs), 0)

     # Example 2: Applying a per-channel-scale
     >>> def scale_channels(input, scale):


def scale_channels(input, scale):
         >>>    scale = scale.refine_names('C')
         >>>    return input * scale.align_as(input)

     >>> num_channels = 3
     >>> scale = torch.randn(num_channels, names=('C',))
     >>> imgs = torch.rand(32, 128, 128, num_channels, names=('N', 'H', 'W', 'C'))
     >>> more_imgs = torch.rand(32, num_channels, 128, 128, names=('N', 'C', 'H', 'W'))
     >>> videos = torch.randn(3, num_channels, 128, 128, 128, names=('N', 'C', 'H', 'W', 'D'))

     # scale_channels is agnostic to the dimension order of the input
     >>> scale_channels(imgs, scale)
     >>> scale_channels(more_imgs, scale)
     >>> scale_channels(videos, scale)

  Warning:

    The named tensor API is experimental and subject to change.

align_to(*names)
  Permutes the dimensions of the "self" tensor to match the order
  specified in "names", adding size-one dims for any new names.


All of the dims of "self" must be named in order to use this
      method. The resulting tensor is a view on the original tensor.
  All dimension names of "self" must be present in "names".
  "names" may contain additional names that are not in
  "self.names"; the output tensor has a size-one dimension for
  each of those new names.

  "names" may contain up to one Ellipsis ("..."). The Ellipsis is
  expanded to be equal to all dimension names of "self" that are
  not mentioned in "names", in the order that they appear in
  "self".

  Python 2 does not support Ellipsis but one may use a string
  literal instead ("'...'").

  Parameters:
     **names** (*iterable of str*) -- The desired dimension
     ordering of the output tensor. May contain up to one Ellipsis
     that is expanded to all unmentioned dim names of "self".

  Examples:

     >>> tensor = torch.randn(2, 2, 2, 2, 2, 2)


named_tensor = tensor.refine_names('A', 'B', 'C', 'D', 'E', 'F')

     # Move the F and E dims to the front while keeping the rest in order
     >>> named_tensor.align_to('F', 'E', ...)

  Warning:

    The named tensor API is experimental and subject to change.

flatten(dims, out_dim) -> Tensor
  Flattens "dims" into a single dimension with name "out_dim".

  All of *dims* must be consecutive in order in the "self" tensor,
  but not necessary contiguous in memory.

  Examples:

     >>> imgs = torch.randn(32, 3, 128, 128, names=('N', 'C', 'H', 'W'))
     >>> flat_imgs = imgs.flatten(['C', 'H', 'W'], 'features')
     >>> flat_imgs.names, flat_imgs.shape
     (('N', 'features'), torch.Size([32, 49152]))

  Warning:

    The named tensor API is experimental and subject to change.


torch.futures
This package provides a "Future" type that encapsulates an
asynchronous execution and a set of utility functions to simplify
operations on "Future" objects. Currently, the "Future" type is
primarily used by the Distributed RPC Framework.
class torch.futures.Future(*, devices=None)
Wrapper around a "torch._C.Future" which encapsulates an
   asynchronous execution of a callable, e.g. "rpc_async()". It also
   exposes a set of APIs to add callback functions and set results.
Warning:
 GPU support is a beta feature, subject to changes.

add_done_callback(callback)
  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", but the order in which they will
  be executed cannot be guaranteed. The callback must take one
  argument, which is the reference to this "Future". The callback


function can use the "value()" method to get the value. Note
      that if this "Future" is already completed, the given callback
      will be run inline.
  We recommend that you use the "then()" method as it provides a
  way to synchronize after your callback has completed.
  "add_done_callback" can be cheaper if your callback does not
  return anything. But both "then()" and "add_done_callback" use
  the same callback registration API under the hood.

  With respect to GPU tensors, this method behaves in the same way
  as "then()".

  Parameters:
     **callback** ("Future") -- a "Callable" that takes in one
     argument, which is the reference to this "Future".

  Note:

    Note that if the callback function throws, either through the
    original future being completed with an exception and calling
    "fut.wait()", or through other code in the callback, error
    handling must be carefully taken care of. For example, if this


callback later completes additional futures, those futures are
        not marked as completed with an error and the user is
        responsible for handling completion/waiting on those futures
        independently.
  Example::
     >>> def callback(fut):
     ...     print("This will run after the future has finished.")
     ...     print(fut.wait())
     >>> fut = torch.futures.Future()
     >>> fut.add_done_callback(callback)
     >>> fut.set_result(5)
     This will run after the future has finished.
     5

done()
  Return "True" if this "Future" is done. A "Future" is done if it
  has a result or an exception.

  If the value contains tensors that reside on GPUs,
  "Future.done()" will return "True" even if the asynchronous
  kernels that are populating those tensors haven't yet completed
  running on the device, because at such stage the result is
  already usable, provided one performs the appropriate


synchronizations (see "wait()").
  Return type:
     bool

set_exception(result)
  Set an exception for this "Future", which will mark this
  "Future" as completed with an error and trigger all attached
  callbacks. Note that when calling wait()/value() on this
  "Future", the exception set here will be raised inline.

  Parameters:
     **result** (*BaseException*) -- the exception for this
     "Future".

  Example::
     >>> fut = torch.futures.Future()
     >>> fut.set_exception(ValueError("foo"))
     >>> fut.wait()
     Traceback (most recent call last):
     ...
     ValueError: foo

set_result(result)
  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.

  If the result contains tensors that reside on GPUs, this method
  can be called even if the asynchronous kernels that are


populating those tensors haven't yet completed running on the
      device, provided that the streams on which those kernels were
      enqueued are set as the current ones when this method is called.
      Put simply, it's safe to call this method immediately after
      launching those kernels, without any additional synchronization,
      as long as one doesn't change streams in between. This method
      will record events on all the relevant current streams and will
      use them to ensure proper scheduling for all the consumers of
      this "Future".
  Parameters:
     **result** (*object*) -- the result object of this "Future".

  Example::
     >>> import threading
     >>> import time
     >>> 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)


...     args=(fut, torch.ones(2) * 3)
         ... )
         >>> t.start()
         >>> print(fut.wait())
         tensor([3., 3.])
         >>> t.join()
then(callback)
  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", but the order in which they will
  be executed cannot be guaranteed (to enforce a certain order
  consider chaining: "fut.then(cb1).then(cb2)"). The callback must
  take one argument, which is the reference to this "Future". The
  callback function can use the "value()" method to get the value.
  Note that if this "Future" is already completed, the given
  callback will be run immediately inline.

  If the "Future"'s value contains tensors that reside on GPUs,
  the callback might be invoked while the async kernels that are
  populating those tensors haven't yet finished executing on the


device. However, the callback will be invoked with some
      dedicated streams set as current (fetched from a global pool)
      which will be synchronized with those kernels. Hence any
      operation performed by the callback on these tensors will be
      scheduled on the device after the kernels complete. In other
      words, as long as the callback doesn't switch streams, it can
      safely manipulate the result without any additional
      synchronization. This is similar to the non-blocking behavior of
      "wait()".
  Similarly, if the callback returns a value that contains tensors
  that reside on a GPU, it can do so even if the kernels that are
  producing these tensors are still running on the device, as long
  as the callback didn't change streams during its execution. If
  one wants to change streams, one must be careful to re-
  synchronize them with the original streams, that is, those that
  were current when the callback was invoked.

  Parameters:


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.

  Return type:
     *Future*[*S*]

  Note:

    Note that if the callback function throws, either through the
    original future being completed with an exception and calling
    "fut.wait()", or through other code in the callback, the
    future returned by "then" will be marked appropriately with
    the encountered error. However, if this callback later
    completes additional futures, those futures are not marked as
    completed with an error and the user is responsible for
    handling completion/waiting on those futures independently.

  Example::
     >>> 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)
         RPC return value is 5.
         Chained cb done. None

value()
  Obtain the value of an already-completed future.

  This method should only be called after a call to "wait()" has
  completed, or inside a callback function passed to "then()". In
  other cases this "Future" may not yet hold a value and calling
  "value()" could fail.

  If the value contains tensors that reside on GPUs, then this
  method will *not* perform any additional synchronization. This
  should be done beforehand, separately, through a call to
  "wait()" (except within callbacks, for which it's already being


taken care of by "then()").
  Returns:
     The value held by this "Future". If the function (callback or
     RPC) creating the value has thrown an error, this "value()"
     method will also throw an error.

  Return type:
     *T*

wait()
  Block until the value of this "Future" is ready.

  If the value contains tensors that reside on GPUs, then an
  additional synchronization is performed with the kernels
  (executing on the device) which may be asynchronously populating
  those tensors. Such sync is non-blocking, which means that
  "wait()" will insert the necessary instructions in the current
  streams to ensure that further operations enqueued on those
  streams will be properly scheduled after the async kernels but,
  once that is done, "wait()" will return, even if those kernels
  are still running. No further synchronization is required when
  accessing and using the values, as long as one doesn't change
  streams.


streams.
  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.

  Return type:
     *T*

torch.futures.collect_all(futures)
Collects the provided "Future" objects into a single combined
   "Future" that is completed when all of the sub-futures are
   completed.
Parameters:
      futures (list) -- a list of "Future" objects.
Returns:
      Returns a "Future" object to a list of the passed in Futures.
Return type:
      Future[List[Future]]
Example::
      >>> fut0 = torch.futures.Future()
      >>> fut1 = torch.futures.Future()
      >>> fut = torch.futures.collect_all([fut0, fut1])
      >>> fut0.set_result(0)
      >>> fut1.set_result(1)
      >>> fut_list = fut.wait()
      >>> print(f"fut0 result = {fut_list[0].wait()}")
      fut0 result = 0
      >>> print(f"fut1 result = {fut_list[1].wait()}")
      fut1 result = 1


fut1 result = 1
torch.futures.wait_all(futures)
Waits for all provided futures to be complete, and returns the list
   of completed values. If any of the futures encounters an error, the
   method will exit early and report the error not waiting for other
   futures to complete.
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.
Return type:
      List


torch.config
torch.config.show()
Return a human-readable string with descriptions of the
   configuration of PyTorch.
torch.config.parallel_info()
Returns detailed string with parallelization settings


torch.profiler
Overview
PyTorch Profiler is a tool that allows the collection of performance
metrics during training and inference. Profiler's context manager API
can be used to better understand what model operators are the most
expensive, examine their input shapes and stack traces, study device
kernel activity and visualize the execution trace.
Note:
An earlier version of the API in "torch.autograd" module is
  considered legacy and will be deprecated.
API Reference
class torch.profiler._KinetoProfile(*, activities=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None)
Low-level profiler wrap the autograd profile
Parameters:
      * activities (iterable) -- list of activity groups (CPU,
        CUDA) to use in profiling, supported values:
        "torch.profiler.ProfilerActivity.CPU",
        "torch.profiler.ProfilerActivity.CUDA". Default value:


ProfilerActivity.CPU and (when available)
        ProfilerActivity.CUDA.
  * **record_shapes** (*bool*) -- save information about
    operator's input shapes.

  * **profile_memory** (*bool*) -- track tensor memory
    allocation/deallocation.

  * **with_stack** (*bool*) -- record source information (file and
    line number) for the ops.

  * **with_flops** (*bool*) -- use formula to estimate the FLOPS
    of specific operators (matrix multiplication and 2D
    convolution).

  * **with_modules** (*bool*) -- record module hierarchy
    (including function names) corresponding to the callstack of
    the op. e.g. If module A's forward call's module B's forward
    which contains an aten::add op, then aten::add's module
    hierarchy is A.B Note that this support exist, at the moment,
    only for TorchScript models and not eager mode models.

  * **experimental_config** (*_ExperimentalConfig*) -- A set of


experimental options used by profiler libraries like Kineto.
        Note, backward compatibility is not guaranteed.
Note:
 This API is experimental and subject to change in the
 future.Enabling shape and stack tracing results in additional
 overhead. When record_shapes=True is specified, profiler will
 temporarily hold references to the tensors; that may further
 prevent certain optimizations that depend on the reference count
 and introduce extra tensor copies.

add_metadata(key, value)
  Adds a user defined metadata with a string key and a string
  value into the trace file

add_metadata_json(key, value)
  Adds a user defined metadata with a string key and a valid json
  value into the trace file

events()
  Returns the list of unaggregated profiler events, to be used in
  the trace callback or after the profiling is finished

export_chrome_trace(path)
  Exports the collected trace in Chrome JSON format.


export_stacks(path, metric='self_cpu_time_total')
  Save stack traces in a file in a format suitable for
  visualization.

  Parameters:
     * **path** (*str*) -- save stacks file to this location;

     * **metric** (*str*) -- metric to use: "self_cpu_time_total"
       or "self_cuda_time_total"

  Note:

    Example of using FlameGraph tool:

    * git clone https://github.com/brendangregg/FlameGraph

    * cd FlameGraph

    * ./flamegraph.pl --title "CPU time" --countname "us."
      profiler.stacks > perf_viz.svg

key_averages(group_by_input_shape=False, group_by_stack_n=0)
  Averages events, grouping them by operator name and (optionally)
  input shapes and stack.

  Note:

    To use shape/stack functionality make sure to set
    record_shapes/with_stack when creating profiler context
    manager.


manager.
class torch.profiler.profile(*, activities=None, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None, use_cuda=None)
Profiler context manager.
Parameters:
      * activities (iterable) -- list of activity groups (CPU,
        CUDA) to use in profiling, supported values:
        "torch.profiler.ProfilerActivity.CPU",
        "torch.profiler.ProfilerActivity.CUDA". Default value:
        ProfilerActivity.CPU and (when available)
        ProfilerActivity.CUDA.
  * **schedule** (*Callable*) -- callable that takes step (int) as
    a single parameter and returns "ProfilerAction" value that
    specifies the profiler action to perform at each step.

  * **on_trace_ready** (*Callable*) -- callable that is called at
    each step when "schedule" returns
    "ProfilerAction.RECORD_AND_SAVE" during the profiling.


record_shapes (bool) -- save information about
        operator's input shapes.

profile_memory (bool) -- track tensor memory
    allocation/deallocation.

with_stack (bool) -- record source information (file and
    line number) for the ops.

with_flops (bool) -- use formula to estimate the FLOPs
    (floating point operations) of specific operators (matrix
    multiplication and 2D convolution).

with_modules (bool) -- record module hierarchy
    (including function names) corresponding to the callstack of
    the op. e.g. If module A's forward call's module B's forward
    which contains an aten::add op, then aten::add's module
    hierarchy is A.B Note that this support exist, at the moment,
    only for TorchScript models and not eager mode models.

experimental_config (_ExperimentalConfig) -- A set of
    experimental options used for Kineto library features. Note,


backward compatibility is not guaranteed.
  * **use_cuda** (*bool*) --

    Deprecated since version 1.8.1: use "activities" instead.

Note:
 Use "schedule()" to generate the callable schedule. Non-default
 schedules are useful when profiling long training jobs and allow
 the user to obtain multiple traces at the different iterations of
 the training process. The default schedule simply records all the
 events continuously for the duration of the context manager.

Note:
 Use "tensorboard_trace_handler()" to generate result files for T
 ensorBoard:"on_trace_ready=torch.profiler.tensorboard_trace_hand
 ler(dir_name)"After profiling, result files can be found in the
 specified directory. Use the command:"tensorboard --logdir
 dir_name"to see the results in TensorBoard. For more information,
 see PyTorch Profiler TensorBoard Plugin

Note:
 Enabling shape and stack tracing results in additional overhead.


When record_shapes=True is specified, profiler will temporarily
     hold references to the tensors; that may further prevent certain
     optimizations that depend on the reference count and introduce
     extra tensor copies.
Examples:
  with torch.profiler.profile(
      activities=[
          torch.profiler.ProfilerActivity.CPU,
          torch.profiler.ProfilerActivity.CUDA,
      ]
  ) as p:
      code_to_profile()
  print(p.key_averages().table(
      sort_by="self_cuda_time_total", row_limit=-1))

Using the profiler's "schedule", "on_trace_ready" and "step"
   functions:
  # Non-default profiler schedule allows user to turn profiler on and off
  # on different iterations of the training loop;
  # trace_handler is called every time a new trace becomes available
  def trace_handler(prof):
      print(prof.key_averages().table(
          sort_by="self_cuda_time_total", row_limit=-1))


prof.export_chrome_trace("/tmp/test_trace_" + str(prof.step_num) + ".json")
  with torch.profiler.profile(
      activities=[
          torch.profiler.ProfilerActivity.CPU,
          torch.profiler.ProfilerActivity.CUDA,
      ],

      # In this example with wait=1, warmup=1, active=2,
      # profiler will skip the first step/iteration,
      # start warming up on the second, record
      # the third and the forth iterations,
      # after which the trace will become available
      # and on_trace_ready (when set) is called;
      # the cycle repeats starting with the next step

      schedule=torch.profiler.schedule(
          wait=1,
          warmup=1,
          active=2),
      on_trace_ready=trace_handler
      # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')
      # used when outputting for tensorboard
      ) as p:
          for iter in range(N):


for iter in range(N):
                  code_iteration_to_profile(iter)
                  # send a signal to the profiler that the next iteration has started
                  p.step()
step()
  Signals the profiler that the next profiling step has started.

class torch.profiler.ProfilerAction(value)
Profiler actions that can be taken at the specified intervals
class torch.profiler.ProfilerActivity
Members:
CPU
CUDA
property name
torch.profiler.schedule(*, wait, warmup, active, repeat=0, skip_first=0)
Returns a callable that can be used as profiler "schedule"
   argument. The profiler will skip the first "skip_first" steps, then
   wait for "wait" steps, then do the warmup for the next "warmup"
   steps, then do the active recording for the next "active" steps and
   then repeat the cycle starting with "wait" steps. The optional
   number of cycles is specified with the "repeat" parameter, the zero


value means that the cycles will continue until the profiling is
   finished.
Return type:
      Callable
torch.profiler.tensorboard_trace_handler(dir_name, worker_name=None, use_gzip=False)
Outputs tracing files to directory of "dir_name", then that
   directory can be directly delivered to tensorboard as logdir.
   "worker_name" should be unique for each worker in distributed
   scenario, it will be set to '[hostname]_[pid]' by default.
Intel Instrumentation and Tracing Technology APIs
torch.profiler.itt.is_available()
Check if ITT feature is available or not
torch.profiler.itt.mark(msg)
Describe an instantaneous event that occurred at some point.
Parameters:
      msg (str) -- ASCII message to associate with the event.
torch.profiler.itt.range_push(msg)
Pushes a range onto a stack of nested range span.  Returns zero-
   based depth of the range that is started.
Parameters:


Parameters:
      msg (str) -- ASCII message to associate with range
torch.profiler.itt.range_pop()
Pops a range off of a stack of nested range spans. Returns the
   zero-based depth of the range that is ended.


Distributed RPC Framework
The distributed RPC framework provides mechanisms for multi-machine
model training through a set of primitives to allow for remote
communication, and a higher-level API to automatically differentiate
models split across several machines.
Warning:
APIs in the RPC package are stable. There are multiple ongoing work
  items to improve performance and error handling, which will ship in
  future releases.
Warning:
CUDA support was introduced in PyTorch 1.9 and is still a beta
  feature. Not all features of the RPC package are yet compatible with
  CUDA support and thus their use is discouraged. These unsupported
  features include: RRefs, JIT compatibility, dist autograd and dist
  optimizer, and profiling. These shortcomings will be addressed in
  future releases.
Note:
Please refer to PyTorch Distributed Overview for a brief
  introduction to all features related to distributed training.
Basics


Basics
The distributed RPC framework makes it easy to run functions remotely,
supports referencing remote objects without copying the real data
around, and provides autograd and optimizer APIs to transparently run
backward and update parameters across RPC boundaries. These features
can be categorized into four sets of APIs.

Remote Procedure Call (RPC) supports running a function on the
   specified destination worker with the given arguments and getting
   the return value back or creating a reference to the return value.
   There are three main RPC APIs: "rpc_sync()" (synchronous),
   "rpc_async()" (asynchronous), and "remote()" (asynchronous and
   returns a reference to the remote return value). Use the
   synchronous API if the user code cannot proceed without the return
   value. Otherwise, use the asynchronous API to get a future, and
   wait on the future when the return value is needed on the caller.
   The "remote()" API is useful when the requirement is to create


something remotely but never need to fetch it to the caller.
   Imagine the case that a driver process is setting up a parameter
   server and a trainer. The driver can create an embedding table on
   the parameter server and then share the reference to the embedding
   table with the trainer, but itself will never use the embedding
   table locally. In this case, "rpc_sync()" and "rpc_async()" are no
   longer appropriate, as they always imply that the return value will
   be returned to the caller immediately or in the future.

Remote Reference (RRef) serves as a distributed shared pointer
   to a local or remote object. It can be shared with other workers
   and reference counting will be handled transparently. Each RRef
   only has one owner and the object only lives on that owner. Non-
   owner workers holding RRefs can get copies of the object from the
   owner by explicitly requesting it. This is useful when a worker
   needs to access some data object, but itself is neither the creator


(the caller of "remote()") or the owner of the object. The
   distributed optimizer, as we will discuss below, is one example of
   such use cases.

Distributed Autograd stitches together local autograd engines
   on all the workers involved in the forward pass, and automatically
   reach out to them during the backward pass to compute gradients.
   This is especially helpful if the forward pass needs to span
   multiple machines when conducting, e.g., distributed model parallel
   training, parameter-server training, etc. With this feature, user
   code no longer needs to worry about how to send gradients across
   RPC boundaries and in which order should the local autograd engines
   be launched, which can become quite complicated where there are
   nested and inter-dependent RPC calls in the forward pass.

Distributed Optimizer's constructor takes a "Optimizer()"
   (e.g., "SGD()", "Adagrad()", etc.) and a list of parameter RRefs,


creates an "Optimizer()" instance on each distinct RRef owner, and
   updates parameters accordingly when running "step()". When you have
   distributed forward and backward passes, parameters and gradients
   will be scattered across multiple workers, and hence it requires an
   optimizer on each of the involved workers. Distributed Optimizer
   wraps all those local optimizers into one, and provides a concise
   constructor and "step()" API.
RPC
Before using RPC and distributed autograd primitives, initialization
must take place. To initialize the RPC framework we need to use
"init_rpc()" which would initialize the RPC framework, RRef framework
and distributed autograd.
torch.distributed.rpc.init_rpc(name, backend=None, rank=- 1, world_size=None, rpc_backend_options=None)
Initializes RPC primitives such as the local RPC agent and
   distributed autograd, which immediately makes the current process
   ready to send and receive RPCs.
Parameters:


ready to send and receive RPCs.
Parameters:
      * name (str) -- a globally unique name of this node.
        (e.g., "Trainer3", "ParameterServer2", "Master", "Worker1")
        Name can only contain number, alphabet, underscore, colon,
        and/or dash, and must be shorter than 128 characters.
  * **backend** (*BackendType**, **optional*) -- The type of RPC
    backend implementation. Supported values is
    "BackendType.TENSORPIPE" (the default). See Backends for more
    information.

  * **rank** (*int*) -- a globally unique id/rank of this node.

  * **world_size** (*int*) -- The number of workers in the group.

  * **rpc_backend_options** (*RpcBackendOptions**, **optional*) --
    The options passed to the RpcAgent constructor. It must be an
    agent-specific subclass of "RpcBackendOptions" and contains
    agent-specific initialization configurations. By default, for
    all agents, it sets the default timeout to 60 seconds and


performs the rendezvous with an underlying process group
        initialized using "init_method = "env://"", meaning that
        environment variables "MASTER_ADDR" and "MASTER_PORT" need to
        be set properly. See Backends for more information and find
        which options are available.
The following APIs allow users to remotely execute functions as well
as create references (RRefs) to remote data objects. In these APIs,
when passing a "Tensor" as an argument or a return value, the
destination worker will try to create a "Tensor" with the same meta
(i.e., shape, stride, etc.). We intentionally disallow transmitting
CUDA tensors because it might crash if the device lists on source and
destination workers do not match. In such cases, applications can
always explicitly move the input tensors to CPU on the caller and move
it to the desired devices on the callee if necessary.
Warning:
TorchScript support in RPC is a prototype feature and subject to


change. Since v1.5.0, "torch.distributed.rpc" supports calling
  TorchScript functions as RPC target functions, and this will help
  improve parallelism on the callee side as executing TorchScript
  functions does not require GIL.
torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None, timeout=- 1.0)
Make a blocking RPC call to run function "func" on worker "to". RPC
   messages are sent and received in parallel to execution of Python
   code. This method is thread-safe.
Parameters:
      * to (str or WorkerInfo or int) --
        name/rank/"WorkerInfo" of the destination worker.
  * **func** (*Callable*) -- a callable function, such as Python
    callables, builtin operators (e.g. "add()") and annotated
    TorchScript functions.

  * **args** (*tuple*) -- the argument tuple for the "func"
    invocation.

  * **kwargs** (*dict*) -- is a dictionary of keyword arguments
    for the "func" invocation.


for the "func" invocation.
  * **timeout** (*float**, **optional*) -- timeout in seconds to
    use for this RPC. If the RPC does not complete in this amount
    of time, an exception indicating it has timed out will be
    raised. A value of 0 indicates an infinite timeout, i.e. a
    timeout error will never be raised. If not provided, the
    default value set during initialization or with
    "_set_rpc_timeout" is used.

Returns:
      Returns the result of running "func" with "args" and "kwargs".
Example::
      Make sure that "MASTER_ADDR" and "MASTER_PORT" are set properly
      on both workers. Refer to "init_process_group()" API for more
      details. For example,
  export MASTER_ADDR=localhost export MASTER_PORT=5678

  Then run the following code in two different processes:

  >>> # On worker 0:
  >>> import torch
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)


ret = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(2), 3))
      >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)
  >>> rpc.shutdown()

  Below is an example of running a TorchScript function using RPC.

  >>> # On both workers:
  >>> @torch.jit.script
  >>> def my_script_add(t1, t2):
  >>>    return torch.add(t1, t2)

  >>> # On worker 0:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)
  >>> ret = rpc.rpc_sync("worker1", my_script_add, args=(torch.ones(2), 3))
  >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)
  >>> rpc.shutdown()

torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None, timeout=- 1.0)
Make a non-blocking RPC call to run function "func" on worker "to".


RPC messages are sent and received in parallel to execution of
   Python code. This method is thread-safe. This method will
   immediately return a "Future" that can be awaited on.
Parameters:
      * to (str or WorkerInfo or int) --
        name/rank/"WorkerInfo" of the destination worker.
  * **func** (*Callable*) -- a callable function, such as Python
    callables, builtin operators (e.g. "add()") and annotated
    TorchScript functions.

  * **args** (*tuple*) -- the argument tuple for the "func"
    invocation.

  * **kwargs** (*dict*) -- is a dictionary of keyword arguments
    for the "func" invocation.

  * **timeout** (*float**, **optional*) -- timeout in seconds to
    use for this RPC. If the RPC does not complete in this amount
    of time, an exception indicating it has timed out will be
    raised. A value of 0 indicates an infinite timeout, i.e. a
    timeout error will never be raised. If not provided, the


default value set during initialization or with
        "_set_rpc_timeout" is used.
Returns:
      Returns a "Future" object that can be waited on. When completed,
      the return value of "func" on "args" and "kwargs" can be
      retrieved from the "Future" object.
Warning:
 Using GPU tensors as arguments or return values of "func" is not
 supported since we don't support sending GPU tensors over the
 wire. You need to explicitly copy GPU tensors to CPU before using
 them as arguments or return values of "func".

Warning:
 The "rpc_async" API does not copy storages of argument tensors
 until sending them over the wire, which could be done by a
 different thread depending on the RPC backend type. The caller
 should make sure that the contents of those tensors stay intact
 until the returned "Future" completes.

Example::
      Make sure that "MASTER_ADDR" and "MASTER_PORT" are set properly


on both workers. Refer to "init_process_group()" API for more
      details. For example,
  export MASTER_ADDR=localhost export MASTER_PORT=5678

  Then run the following code in two different processes:

  >>> # On worker 0:
  >>> import torch
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)
  >>> fut1 = rpc.rpc_async("worker1", torch.add, args=(torch.ones(2), 3))
  >>> fut2 = rpc.rpc_async("worker1", min, args=(1, 2))
  >>> result = fut1.wait() + fut2.wait()
  >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)
  >>> rpc.shutdown()

  Below is an example of running a TorchScript function using RPC.

  >>> # On both workers:
  >>> @torch.jit.script
  >>> def my_script_add(t1, t2):
  >>>    return torch.add(t1, t2)

  >>> # On worker 0:
  >>> import torch.distributed.rpc as rpc


import torch.distributed.rpc as rpc
      >>> rpc.init_rpc("worker0", rank=0, world_size=2)
      >>> fut = rpc.rpc_async("worker1", my_script_add, args=(torch.ones(2), 3))
      >>> ret = fut.wait()
      >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)
  >>> rpc.shutdown()

torch.distributed.rpc.remote(to, func, args=None, kwargs=None, timeout=- 1.0)
Make a remote call to run "func" on worker "to" and return an
   "RRef" to the result value immediately. Worker "to" will be the
   owner of the returned "RRef", and the worker calling "remote" is a
   user. The owner manages the global reference count of its "RRef",
   and the owner "RRef" is only destructed when globally there are no
   living references to it.
Parameters:
      * to (str or WorkerInfo or int) --
        name/rank/"WorkerInfo" of the destination worker.


func (Callable) -- a callable function, such as Python
        callables, builtin operators (e.g. "add()") and annotated
        TorchScript functions.

args (tuple) -- the argument tuple for the "func"
    invocation.

kwargs (dict) -- is a dictionary of keyword arguments
    for the "func" invocation.

timeout (float, optional) -- timeout in seconds for
    this remote call. If the creation of this "RRef" on worker
    "to" is not successfully processed on this worker within this
    timeout, then the next time there is an attempt to use the
    RRef (such as "to_here()"), a timeout will be raised
    indicating this failure. A value of 0 indicates an infinite
    timeout, i.e. a timeout error will never be raised. If not
    provided, the default value set during initialization or with
    "_set_rpc_timeout" is used.

Returns:
      A user "RRef" instance to the result value. Use the blocking API


"torch.distributed.rpc.RRef.to_here()" to retrieve the result
      value locally.
Warning:
 The "remote" API does not copy storages of argument tensors until
 sending them over the wire, which could be done by a different
 thread depending on the RPC backend type. The caller should make
 sure that the contents of those tensors stay intact until the
 returned RRef is confirmed by the owner, which can be checked
 using the "torch.distributed.rpc.RRef.confirmed_by_owner()" API.

Warning:
 Errors such as timeouts for the "remote" API are handled on a
 best-effort basis. This means that when remote calls initiated by
 "remote" fail, such as with a timeout error, we take a best-
 effort approach to error handling. This means that errors are
 handled and set on the resulting RRef on an asynchronous basis.
 If the RRef has not been used by the application before this
 handling (such as "to_here" or fork call), then future uses of


the "RRef" will appropriately raise errors. However, it is
     possible that the user application will use the "RRef" before the
     errors are handled. In this case, errors may not be raised as
     they have not yet been handled.
Example:
  Make sure that ``MASTER_ADDR`` and ``MASTER_PORT`` are set properly
  on both workers. Refer to :meth:`~torch.distributed.init_process_group`
  API for more details. For example,

  export MASTER_ADDR=localhost
  export MASTER_PORT=5678

  Then run the following code in two different processes:

  >>> # On worker 0:
  >>> import torch
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)
  >>> rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
  >>> rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
  >>> x = rref1.to_here() + rref2.to_here()
  >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc


import torch.distributed.rpc as rpc
      >>> rpc.init_rpc("worker1", rank=1, world_size=2)
      >>> rpc.shutdown()

  Below is an example of running a TorchScript function using RPC.

  >>> # On both workers:
  >>> @torch.jit.script
  >>> def my_script_add(t1, t2):
  >>>    return torch.add(t1, t2)

  >>> # On worker 0:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)
  >>> rref = rpc.remote("worker1", my_script_add, args=(torch.ones(2), 3))
  >>> rref.to_here()
  >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)
  >>> rpc.shutdown()

torch.distributed.rpc.get_worker_info(worker_name=None)
Get "WorkerInfo" of a given worker name. Use this "WorkerInfo" to
   avoid passing an expensive string on every invocation.
Parameters:
      worker_name (str) -- the string name of a worker. If


"None", return the the id of the current worker. (default
      "None")
Returns:
      "WorkerInfo" instance for the given "worker_name" or
      "WorkerInfo" of the current worker if "worker_name" is "None".
torch.distributed.rpc.shutdown(graceful=True, timeout=0)
Perform a shutdown of the RPC agent, and then destroy the RPC
   agent. This stops the local agent from accepting outstanding
   requests, and shuts down the RPC framework by terminating all RPC
   threads. If "graceful=True", this will block until all local and
   remote RPC processes reach this method and wait for all outstanding
   work to complete. Otherwise, if "graceful=False", this is a local
   shutdown, and it does not wait for other RPC processes to reach
   this method.
Warning:
 For "Future" objects returned by "rpc_async()", "future.wait()"
 should not be called after "shutdown()".

Parameters:
      graceful (bool) -- Whether to do a graceful shutdown or


not. If True, this will 1) wait until there is no pending system
      messages for "UserRRefs" and delete them; 2) block until all
      local and remote RPC processes have reached this method and wait
      for all outstanding work to complete.
Example::
      Make sure that "MASTER_ADDR" and "MASTER_PORT" are set properly
      on both workers. Refer to "init_process_group()" API for more
      details. For example,
  export MASTER_ADDR=localhost export MASTER_PORT=5678

  Then run the following code in two different processes:

  >>> # On worker 0:
  >>> import torch
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)
  >>> # do some work
  >>> result = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(1), 1))
  >>> # ready to shutdown
  >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch.distributed.rpc as rpc
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)


wait for worker 0 to finish work, and then shutdown.
  >>> rpc.shutdown()

class torch.distributed.rpc.WorkerInfo
A structure that encapsulates information of a worker in the
   system. Contains the name and ID of the worker. This class is not
   meant to be constructed directly, rather, an instance can be
   retrieved through "get_worker_info()" and the result can be passed
   in to functions such as "rpc_sync()", "rpc_async()", "remote()" to
   avoid copying a string on every invocation.
property id
  Globally unique id to identify the worker.

property name
  The name of the worker.

The RPC package also provides decorators which allow applications to
specify how a given function should be treated on the callee side.
torch.distributed.rpc.functions.async_execution(fn)
A decorator for a function indicating that the return value of the
   function is guaranteed to be a "Future" object and this function


can run asynchronously on the RPC callee. More specifically, the
   callee extracts the "Future" returned by the wrapped function and
   installs subsequent processing steps as a callback to that
   "Future". The installed callback will read the value from the
   "Future" when completed and send the value back as the RPC
   response. That also means the returned "Future" only exists on the
   callee side and is never sent through RPC. This decorator is useful
   when the wrapped function's ("fn") execution needs to pause and
   resume due to, e.g., containing "rpc_async()" or waiting for other
   signals.
Note:
 To enable asynchronous execution, applications must pass the
 function object returned by this decorator to RPC APIs. If RPC
 detected attributes installed by this decorator, it knows that
 this function returns a "Future" object and will handle that
 accordingly. However, this does not mean this decorator has to be


outmost one when defining a function. For example, when combined
     with "@staticmethod" or "@classmethod",
     "@rpc.functions.async_execution" needs to be the inner decorator
     to allow the target function be recognized as a static or class
     function. This target function can still execute asynchronously
     because, when accessed, the static or class method preserves
     attributes installed by "@rpc.functions.async_execution".
Example::
      The returned "Future" object can come from "rpc_async()",
      "then()", or "Future" constructor. The example below shows
      directly using the "Future" returned by "then()".
  >>> from torch.distributed import rpc
  >>>
  >>> # omitting setup and shutdown RPC
  >>>
  >>> # On all workers
  >>> @rpc.functions.async_execution
  >>> def async_add_chained(to, x, y, z):
  >>>     # This function runs on "worker1" and returns immediately when


# the callback is installed through the `then(cb)` API. In the
  >>>     # mean time, the `rpc_async` to "worker2" can run concurrently.
  >>>     # When the return value of that `rpc_async` arrives at
  >>>     # "worker1", "worker1" will run the lambda function accordingly
  >>>     # and set the value for the previously returned `Future`, which
  >>>     # will then trigger RPC to send the result back to "worker0".
  >>>     return rpc.rpc_async(to, torch.add, args=(x, y)).then(
  >>>         lambda fut: fut.wait() + z
  >>>     )
  >>>
  >>> # On worker0
  >>> ret = rpc.rpc_sync(
  >>>     "worker1",
  >>>     async_add_chained,
  >>>     args=("worker2", torch.ones(2), 1, 1)
  >>> )
  >>> print(ret)  # prints tensor([3., 3.])

  When combined with TorchScript decorators, this decorator must
  be the outmost one.

  >>> from torch import Tensor
  >>> from torch.futures import Future


from torch.futures import Future
      >>> from torch.distributed import rpc
      >>>
      >>> # omitting setup and shutdown RPC
      >>>
      >>> # On all workers
      >>> @torch.jit.script
      >>> def script_add(x: Tensor, y: Tensor) -> Tensor:
      >>>     return x + y
      >>>
      >>> @rpc.functions.async_execution
      >>> @torch.jit.script
      >>> def async_add(to: str, x: Tensor, y: Tensor) -> Future[Tensor]:
      >>>     return rpc.rpc_async(to, script_add, (x, y))
      >>>
      >>> # On worker0
      >>> ret = rpc.rpc_sync(
      >>>     "worker1",
      >>>     async_add,
      >>>     args=("worker2", torch.ones(2), 1)
      >>> )
      >>> print(ret)  # prints tensor([2., 2.])

  When combined with static or class method, this decorator must
  be the inner one.

  >>> from torch.distributed import rpc
  >>>
  >>> # omitting setup and shutdown RPC
  >>>
  >>> # On all workers
  >>> class AsyncExecutionClass:
  >>>


class AsyncExecutionClass:
      >>>
      >>>     @staticmethod
      >>>     @rpc.functions.async_execution
      >>>     def static_async_add(to, x, y, z):
      >>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
      >>>             lambda fut: fut.wait() + z
      >>>         )
      >>>
      >>>     @classmethod
      >>>     @rpc.functions.async_execution
      >>>     def class_async_add(cls, to, x, y, z):
      >>>         ret_fut = torch.futures.Future()
      >>>         rpc.rpc_async(to, torch.add, args=(x, y)).then(
      >>>             lambda fut: ret_fut.set_result(fut.wait() + z)
      >>>         )
      >>>         return ret_fut
      >>>
      >>>     @rpc.functions.async_execution
      >>>     def bound_async_add(self, to, x, y, z):
      >>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
      >>>             lambda fut: fut.wait() + z
      >>>         )
      >>>
      >>> # On worker0
      >>> ret = rpc.rpc_sync(
      >>>     "worker1",


"worker1",
  >>>     AsyncExecutionClass.static_async_add,
  >>>     args=("worker2", torch.ones(2), 1, 2)
  >>> )
  >>> print(ret)  # prints tensor([4., 4.])
  >>>
  >>> ret = rpc.rpc_sync(
  >>>     "worker1",
  >>>     AsyncExecutionClass.class_async_add,
  >>>     args=("worker2", torch.ones(2), 1, 2)
  >>> )
  >>> print(ret)  # prints tensor([4., 4.])

  This decorator also works with RRef helpers, i.e., .
  "torch.distributed.rpc.RRef.rpc_sync()",
  "torch.distributed.rpc.RRef.rpc_async()", and
  "torch.distributed.rpc.RRef.remote()".

  >>> from torch.distributed import rpc
  >>>
  >>> # reuse the AsyncExecutionClass class above
  >>> rref = rpc.remote("worker1", AsyncExecutionClass)
  >>> ret = rref.rpc_sync().static_async_add("worker2", torch.ones(2), 1, 2)
  >>> print(ret)  # prints tensor([4., 4.])
  >>>
  >>> rref = rpc.remote("worker1", AsyncExecutionClass)


ret = rref.rpc_async().static_async_add("worker2", torch.ones(2), 1, 2).wait()
      >>> print(ret)  # prints tensor([4., 4.])
      >>>
      >>> rref = rpc.remote("worker1", AsyncExecutionClass)
      >>> ret = rref.remote().static_async_add("worker2", torch.ones(2), 1, 2).to_here()
      >>> print(ret)  # prints tensor([4., 4.])

Backends
The RPC module can leverage different backends to perform the
communication between the nodes. The backend to be used can be
specified in the "init_rpc()" function, by passing a certain value of
the "BackendType" enum. Regardless of what backend is used, the rest
of the RPC API won't change. Each backend also defines its own
subclass of the "RpcBackendOptions" class, an instance of which can
also be passed to "init_rpc()" to configure the backend's behavior.
class torch.distributed.rpc.BackendType(value)
An enum class of available backends.
PyTorch ships with a builtin "BackendType.TENSORPIPE" backend.


Additional ones can be registered using the "register_backend()"
   function.
class torch.distributed.rpc.RpcBackendOptions
An abstract structure encapsulating the options passed into the RPC
   backend. An instance of this class can be passed in to "init_rpc()"
   in order to initialize RPC with specific configurations, such as
   the RPC timeout and "init_method" to be used.
property init_method
  URL specifying how to initialize the process group. Default is
  "env://"

property rpc_timeout
  A float indicating the timeout to use for all RPCs. If an RPC
  does not complete in this timeframe, it will complete with an
  exception indicating that it has timed out.

TensorPipe Backend
~~~~~~~~~~~~~~~~~~
The TensorPipe agent, which is the default, leverages the TensorPipe
library, which provides a natively point-to-point communication
primitive specifically suited for machine learning that fundamentally
addresses some of the limitations of Gloo. Compared to Gloo, it has


the advantage of being asynchronous, which allows a large number of
transfers to occur simultaneously, each at their own speed, without
blocking each other. It will only open pipes between pairs of nodes
when needed, on demand, and when one node fails only its incident
pipes will be closed, while all other ones will keep working as
normal. In addition, it is able to support multiple different
transports (TCP, of course, but also shared memory, NVLink,
InfiniBand, ...) and can automatically detect their availability and
negotiate the best transport to use for each pipe.
The TensorPipe backend has been introduced in PyTorch v1.6 and is
being actively developed. At the moment, it only supports CPU tensors,
with GPU support coming soon. It comes with a TCP-based transport,
just like Gloo. It is also able to automatically chunk and multiplex
large tensors over multiple sockets and threads in order to achieve
very high bandwidths. The agent will be able to pick the best


transport on its own, with no intervention required.
Example:

import os
from torch.distributed import rpc
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
rpc.init_rpc(
    "worker1",
    rank=0,
    world_size=2,
    rpc_backend_options=rpc.TensorPipeRpcBackendOptions(
        num_worker_threads=8,
        rpc_timeout=20 # 20 second timeout
    )
)
omitting init_rpc invocation on worker2

class torch.distributed.rpc.TensorPipeRpcBackendOptions(*, num_worker_threads=16, rpc_timeout=60.0, init_method='env://', device_maps=None, devices=None, _transports=None, _channels=None)
The backend options for "TensorPipeAgent", derived from
   "RpcBackendOptions".
Parameters:
      * num_worker_threads (int, optional) -- The number of
        threads in the thread-pool used by "TensorPipeAgent" to
        execute requests (default: 16).


execute requests (default: 16).
  * **rpc_timeout** (*float**, **optional*) -- The default
    timeout, in seconds, for RPC requests (default: 60 seconds).
    If the RPC has not completed in this timeframe, an exception
    indicating so will be raised. Callers can override this
    timeout for individual RPCs in "rpc_sync()" and "rpc_async()"
    if necessary.

  * **init_method** (*str**, **optional*) -- The URL to initialize
    the distributed store used for rendezvous. It takes any value
    accepted for the same argument of "init_process_group()"
    (default: "env://").

  * **device_maps** (*Dict**[**str**, **Dict**]**, **optional*) --
    Device placement mappings from this worker to the callee. Key
    is the callee worker name and value the dictionary ("Dict" of
    "int", "str", or "torch.device") that maps this worker's
    devices to the callee worker's devices. (default: "None")


devices (List[int, str, or "torch.device"], optional) --
        all local CUDA devices used by RPC agent. By Default, it will
        be initialized to all local devices from its own "device_maps"
        and corresponding devices from its peers' "device_maps". When
        processing CUDA RPC requests, the agent will properly
        synchronize CUDA streams for all devices in this "List".

property device_maps
  The device map locations.

property devices
  All devices used by the local agent.

property init_method
  URL specifying how to initialize the process group. Default is
  "env://"

property num_worker_threads
  The number of threads in the thread-pool used by
  "TensorPipeAgent" to execute requests.

property rpc_timeout
  A float indicating the timeout to use for all RPCs. If an RPC
  does not complete in this timeframe, it will complete with an
  exception indicating that it has timed out.

set_device_map(to, device_map)


set_device_map(to, device_map)
  Set device mapping between each RPC caller and callee pair. This
  function can be called multiple times to incrementally add
  device placement configurations.

  Parameters:
     * **to** (*str*) -- Callee name.

     * **device_map** (*Dict of python:int**, **str**, or
       **torch.device*) -- Device placement mappings from this
       worker to the callee. This map must be invertible.

  -[ Example ]-

  >>> # both workers
  >>> def add(x, y):
  >>>     print(x)  # tensor([1., 1.], device='cuda:1')
  >>>     return x + y, (x + y).to(2)
  >>>
  >>> # on worker 0
  >>> options = TensorPipeRpcBackendOptions(
  >>>     num_worker_threads=8,
  >>>     device_maps={"worker1": {0: 1}}
  >>>     # maps worker0's cuda:0 to worker1's cuda:1
  >>> )
  >>> options.set_device_map("worker1", {1: 2})
  >>> # maps worker0's cuda:1 to worker1's cuda:2
  >>>
  >>> rpc.init_rpc(


  >>> rpc.init_rpc(
  >>>     "worker0",
  >>>     rank=0,
  >>>     world_size=2,
  >>>     backend=rpc.BackendType.TENSORPIPE,
  >>>     rpc_backend_options=options
  >>> )
  >>>
  >>> x = torch.ones(2)
  >>> rets = rpc.rpc_sync("worker1", add, args=(x.to(0), 1))
  >>> # The first argument will be moved to cuda:1 on worker1. When
  >>> # sending the return value back, it will follow the invert of
  >>> # the device map, and hence will be moved back to cuda:0 and
  >>> # cuda:1 on worker0
  >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
  >>> print(rets[1])  # tensor([2., 2.], device='cuda:1')

set_devices(devices)
  Set local devices used by the TensorPipe RPC agent. When
  processing CUDA RPC requests, the TensorPipe RPC agent will
  properly synchronize CUDA streams for all devices in this
  "List".

  Parameters:
     **devices** (*List of python:int**, **str**, or


*torch.device) -- local devices used by the TensorPipe RPC
         agent.
Note:
The RPC framework does not automatically retry any "rpc_sync()",
  "rpc_async()" and "remote()" calls. The reason being that there is
  no way the RPC framework can determine whether an operation is
  idempotent or not and whether it is safe to retry. As a result, it
  is the application's responsibility to deal with failures and retry
  if necessary. RPC communication is based on TCP and as a result
  failures could happen due to network failures or intermittent
  network connectivity issues. In such scenarios, the application
  needs to retry appropriately with reasonable backoffs to ensure the
  network isn't overwhelmed by aggressive retries.
RRef
Warning:
RRefs are not currently supported when using CUDA tensors
An "RRef" (Remote REFerence) is a reference to a value of some type
"T" (e.g. "Tensor") on a remote worker. This handle keeps the
referenced remote value alive on the owner, but there is no


implication that the value will be transferred to the local worker in
the future. RRefs can be used in multi-machine training by holding
references to nn.Modules that exist on other workers, and calling the
appropriate functions to retrieve or modify their parameters during
training. See Remote Reference Protocol for more details.
class torch.distributed.rpc.RRef
More Information about RRef
^^^^^^^^^^^^^^^^^^^^^^^^^^^

Remote Reference Protocol

Background

Assumptions

RRef Lifetime

Design Reasoning

Implementation

Protocol Scenarios

User Share RRef with Owner as Return Value

User Share RRef with Owner as Argument

Owner Share RRef with User

User Share RRef with User

RemoteModule
Warning:
RemoteModule is not currently supported when using CUDA tensors
"RemoteModule" is an easy way to create an nn.Module remotely on a
different process. The actual module resides on a remote host, but the


local host has a handle to this module and invoke this module similar
to a regular nn.Module. The invocation however incurs RPC calls to the
remote end and can be performed asynchronously if needed via
additional APIs supported by RemoteModule.
class torch.distributed.nn.api.remote_module.RemoteModule(args, *kwargs)
  A RemoteModule instance can only be created after RPC
  initialization. It creates a user-specified module on a
  specified remote node. It behaves like a regular "nn.Module"
  except that the "forward" method is executed on the remote node.
  It takes care of autograd recording to ensure the backward pass
  propagates gradients back to the corresponding remote module.

  It generates two methods "forward_async" and "forward" based on
  the signature of the "forward" method of "module_cls".
  "forward_async" runs asynchronously and returns a Future. The
  arguments of "forward_async" and "forward" are the same as the


"forward" method of the module returned by the "module_cls".
  For example, if "module_cls" returns an instance of "nn.Linear",
  that has "forward" method signature: "def forward(input: Tensor)
  -> Tensor:", the generated "RemoteModule" will have 2 methods
  with the signatures:

     "def forward(input: Tensor) -> Tensor:"
     "def forward_async(input: Tensor) -> Future[Tensor]:"

Parameters:
      * remote_device (str) -- Device on the destination worker
        where we'd like to place this module. The format should be
        "/", where the device field can be parsed
        as torch.device type. E.g., "trainer0/cpu", "trainer0",
        "ps0/cuda:0". In addition, the device field can be optional
        and the default value is "cpu".
  * **module_cls** (*nn.Module*) --

    Class for the module to be created remotely. For example,

    >>> class MyModule(nn.Module):
    >>>     def forward(input):


def forward(input):
    >>>         return input + 1
    >>>
    >>> module_cls = MyModule

  * **args** (*Sequence**, **optional*) -- args to be passed to
    "module_cls".

  * **kwargs** (*Dict**, **optional*) -- kwargs to be passed to
    "module_cls".

Returns:
      A remote module instance which wraps the "Module" created by the
      user-provided "module_cls", it has a blocking "forward" method
      and an asynchronous "forward_async" method that returns a future
      of the "forward" call on the user-provided module on the remote
      side.
Example::
      Run the following code in two different processes:
  >>> # On worker 0:
  >>> import torch
  >>> import torch.distributed.rpc as rpc
  >>> from torch import nn, Tensor
  >>> from torch.distributed.nn.api.remote_module import RemoteModule
  >>>
  >>> rpc.init_rpc("worker0", rank=0, world_size=2)
  >>> remote_linear_module = RemoteModule(


remote_linear_module = RemoteModule(
      >>>     "worker1/cpu", nn.Linear, args=(20, 30),
      >>> )
      >>> input = torch.randn(128, 20)
      >>> ret_fut = remote_linear_module.forward_async(input)
      >>> ret = ret_fut.wait()
      >>> rpc.shutdown()

  >>> # On worker 1:
  >>> import torch
  >>> import torch.distributed.rpc as rpc
  >>>
  >>> rpc.init_rpc("worker1", rank=1, world_size=2)
  >>> rpc.shutdown()

  Furthermore, a more practical example that is combined with
  DistributedDataParallel (DDP) can be found in this tutorial.

get_module_rref()
  Returns an "RRef" ("RRef[nn.Module]") pointing to the remote
  module.

  Return type:
     *RRef*[*Module*]

remote_parameters(recurse=True)
  Returns a list of "RRef" pointing to the remote module's
  parameters. This can typically be used in conjuction with
  "DistributedOptimizer".

  Parameters:
     **recurse** (*bool*) -- if True, then returns parameters of


the remote module and all submodules of the remote module.
         Otherwise, returns only parameters that are direct members of
         the remote module.
  Returns:
     A list of "RRef" ("List[RRef[nn.Parameter]]") to remote
     module's parameters.

  Return type:
     *List*[*RRef*[*Parameter*]]

Distributed Autograd Framework
Warning:
Distributed autograd is not currently supported when using CUDA
  tensors
This module provides an RPC-based distributed autograd framework that
can be used for applications such as model parallel training. In
short, applications may send and receive gradient recording tensors
over RPC. In the forward pass, we record when gradient recording
tensors are sent over RPC and during the backward pass we use this
information to perform a distributed backward pass using RPC. For more
details see Distributed Autograd Design.


details see Distributed Autograd Design.
torch.distributed.autograd.backward(context_id: int, roots: List[Tensor], retain_graph=False) -> None
Kicks off the distributed backward pass using the provided roots.
   This currently implements the FAST mode algorithm which assumes all
   RPC messages sent in the same distributed autograd context across
   workers would be part of the autograd graph during the backward
   pass.
We use the provided roots to discover the autograd graph and
   compute appropriate dependencies. This method blocks until the
   entire autograd computation is done.
We accumulate the gradients in the appropriate
   "torch.distributed.autograd.context" on each of the nodes. The
   autograd context to be used is looked up given the "context_id"
   that is passed in when "torch.distributed.autograd.backward()" is
   called. If there is no valid autograd context corresponding to the
   given ID, we throw an error. You can retrieve the accumulated


gradients using the "get_gradients()" API.
Parameters:
      * context_id (int) -- The autograd context id for which we
        should retrieve the gradients.
  * **roots** (*list*) -- Tensors which represent the roots of the
    autograd computation. All the tensors should be scalars.

  * **retain_graph** (*bool**, **optional*) -- If False, the graph
    used to compute the grad will be freed. Note that in nearly
    all cases setting this option to True is not needed and often
    can be worked around in a much more efficient way. Usually,
    you need to set this to True to run backward multiple times.

Example::
      >>> import torch.distributed.autograd as dist_autograd
      >>> with dist_autograd.context() as context_id:
      >>>     pred = model.forward()
      >>>     loss = loss_func(pred, loss)
      >>>     dist_autograd.backward(context_id, loss)
class torch.distributed.autograd.context


class torch.distributed.autograd.context
Context object to wrap forward and backward passes when using
   distributed autograd. The "context_id" generated in the "with"
   statement  is required to uniquely identify a distributed backward
   pass on all workers. Each worker stores metadata associated with
   this "context_id", which is required to correctly execute a
   distributed autograd pass.
Example::
      >>> import torch.distributed.autograd as dist_autograd
      >>> with dist_autograd.context() as context_id:
      >>>     t1 = torch.rand((3, 3), requires_grad=True)
      >>>     t2 = torch.rand((3, 3), requires_grad=True)
      >>>     loss = rpc.rpc_sync("worker1", torch.add, args=(t1, t2)).sum()
      >>>     dist_autograd.backward(context_id, [loss])
torch.distributed.autograd.get_gradients(context_id: int) -> Dict[Tensor, Tensor]
Retrieves a map from Tensor to the appropriate gradient for that
   Tensor accumulated in the provided context corresponding to the


given "context_id" as part of the distributed autograd backward
   pass.
Parameters:
      context_id (int) -- The autograd context id for which we
      should retrieve the gradients.
Returns:
      A map where the key is the Tensor and the value is the
      associated gradient for that Tensor.
Example::
      >>> import torch.distributed.autograd as dist_autograd
      >>> with dist_autograd.context() as context_id:
      >>>     t1 = torch.rand((3, 3), requires_grad=True)
      >>>     t2 = torch.rand((3, 3), requires_grad=True)
      >>>     loss = t1 + t2
      >>>     dist_autograd.backward(context_id, [loss.sum()])
      >>>     grads = dist_autograd.get_gradients(context_id)
      >>>     print(grads[t1])
      >>>     print(grads[t2])
More Information about RPC Autograd
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Distributed Autograd Design

Background

Autograd recording during the forward pass

Distributed Autograd Context

Distributed Backward Pass


Distributed Backward Pass

Computing dependencies

FAST mode algorithm

SMART mode algorithm

Distributed Optimizer

Simple end to end example

Distributed Optimizer
See the torch.distributed.optim page for documentation on distributed
optimizers.
Design Notes
The distributed autograd design note covers the design of the RPC-
based distributed autograd framework that is useful for applications
such as model parallel training.

Distributed Autograd Design

The RRef design note covers the design of the RRef (Remote REFerence)
protocol used to refer to values on remote workers by the framework.

Remote Reference Protocol

Tutorials
The RPC tutorials introduce users to the RPC framework, provide
several example applications using torch.distributed.rpc APIs, and
demonstrate how to use the profiler to profile RPC-based workloads.

Getting started with Distributed RPC Framework


Getting started with Distributed RPC Framework

Implementing a Parameter Server using Distributed RPC Framework

Combining Distributed DataParallel with Distributed RPC Framework
  (covers RemoteModule as well)

Profiling RPC-based Workloads

Implementing batch RPC processing

Distributed Pipeline Parallel


torch.special
The torch.special module, modeled after SciPy's special module.
Functions
torch.special.airy_ai(input, *, out=None) -> Tensor
Airy function \text{Ai}\left(\text{input}\right).
Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
torch.special.bessel_j0(input, *, out=None) -> Tensor
Bessel function of the first kind of order 0.
Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
torch.special.bessel_j1(input, *, out=None) -> Tensor
Bessel function of the first kind of order 1.
Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
torch.special.digamma(input, *, out=None) -> Tensor
Computes the logarithmic derivative of the gamma function on
   input.


input.
  \digamma(x) = \frac{d}{dx} \ln\left(\Gamma\left(x\right)\right)
  = \frac{\Gamma'(x)}{\Gamma(x)}

Parameters:
      input (Tensor) -- the tensor to compute the digamma
      function on
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Note:
 This function is similar to SciPy's *scipy.special.digamma*.

Note:
 From PyTorch 1.8 onwards, the digamma function returns *-Inf* for
 *0*. Previously it returned *NaN* for *0*.

Example:
  >>> a = torch.tensor([1, 0.5])
  >>> torch.special.digamma(a)
  tensor([-0.5772, -1.9635])

torch.special.entr(input, *, out=None) -> Tensor
Computes the entropy on "input" (as defined below), elementwise.
  \begin{align} \text{entr(x)} = \begin{cases}     -x * \ln(x)  &
  x > 0 \\     0 &  x = 0.0 \\     -\infty & x < 0 \end{cases}
  \end{align}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:


Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> a = torch.arange(-0.5, 1, 0.5)
      >>> a
      tensor([-0.5000,  0.0000,  0.5000])
      >>> torch.special.entr(a)
      tensor([  -inf, 0.0000, 0.3466])
torch.special.erf(input, *, out=None) -> Tensor
Computes the error function of "input". The error function is
   defined as follows:
  \mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_{0}^{x} e^{-t^2} dt

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.special.erf(torch.tensor([0, -1., 10.]))
  tensor([ 0.0000, -0.8427,  1.0000])

torch.special.erfc(input, *, out=None) -> Tensor
Computes the complementary error function of "input". The
   complementary error function is defined as follows:
  \mathrm{erfc}(x) = 1 - \frac{2}{\sqrt{\pi}} \int_{0}^{x}
  e^{-t^2} dt

Parameters:


e^{-t^2} dt
Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.special.erfc(torch.tensor([0, -1., 10.]))
  tensor([ 1.0000, 1.8427,  0.0000])

torch.special.erfcx(input, *, out=None) -> Tensor
Computes the scaled complementary error function for each element
   of "input". The scaled complementary error function is defined as
   follows:
  \mathrm{erfcx}(x) = e^{x^2} \mathrm{erfc}(x)

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.special.erfcx(torch.tensor([0, -1., 10.]))
  tensor([ 1.0000, 5.0090, 0.0561])

torch.special.erfinv(input, *, out=None) -> Tensor
Computes the inverse error function of "input". The inverse error
   function is defined in the range (-1, 1) as:
  \mathrm{erfinv}(\mathrm{erf}(x)) = x

Parameters:


Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.special.erfinv(torch.tensor([0, 0.5, -1.]))
  tensor([ 0.0000,  0.4769,    -inf])

torch.special.exp2(input, *, out=None) -> Tensor
Computes the base two exponential function of "input".
  y_{i} = 2^{x_{i}}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.special.exp2(torch.tensor([0, math.log2(2.), 3, 4]))
  tensor([ 1.,  2.,  8., 16.])

torch.special.expit(input, *, out=None) -> Tensor
Computes the expit (also known as the logistic sigmoid function) of
   the elements of "input".
  \text{out}_{i} = \frac{1}{1 + e^{-\text{input}_{i}}}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.


Example:
  >>> t = torch.randn(4)
  >>> t
  tensor([ 0.9213,  1.0887, -0.8858, -1.7683])
  >>> torch.special.expit(t)
  tensor([ 0.7153,  0.7481,  0.2920,  0.1458])

torch.special.expm1(input, *, out=None) -> Tensor
Computes the exponential of the elements minus 1 of "input".
  y_{i} = e^{x_{i}} - 1

Note:
 This function provides greater precision than exp(x) - 1 for
 small values of x.

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.special.expm1(torch.tensor([0, math.log(2.)]))
  tensor([ 0.,  1.])

torch.special.gammainc(input, other, *, out=None) -> Tensor
Computes the regularized lower incomplete gamma function:
  \text{out}_{i} = \frac{1}{\Gamma(\text{input}_i)}
  \int_0^{\text{other}_i} t^{\text{input}_i-1} e^{-t} dt

where both \text{input}_i and \text{other}_i are weakly positive


and at least one is strictly positive. If both are zero or either
   is negative then \text{out}_i=\text{nan}. \Gamma(\cdot) in the
   equation above is the gamma function,
  \Gamma(\text{input}_i) = \int_0^\infty t^{(\text{input}_i-1)}
  e^{-t} dt.

See "torch.special.gammaincc()" and "torch.special.gammaln()" for
   related functions.
Supports broadcasting to a common shape and float inputs.
Note:
 The backward pass with respect to "input" is not yet supported.
 Please open an issue on PyTorch's Github to request it.

Parameters:
      * input (Tensor) -- the first non-negative input tensor
  * **other** (*Tensor*) -- the second non-negative input tensor

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a1 = torch.tensor([4.0])
  >>> a2 = torch.tensor([3.0, 4.0, 5.0])
  >>> a = torch.special.gammaincc(a1, a2)
  tensor([0.3528, 0.5665, 0.7350])
  tensor([0.3528, 0.5665, 0.7350])


tensor([0.3528, 0.5665, 0.7350])
      >>> b = torch.special.gammainc(a1, a2) + torch.special.gammaincc(a1, a2)
      tensor([1., 1., 1.])
torch.special.gammaincc(input, other, *, out=None) -> Tensor
Computes the regularized upper incomplete gamma function:
  \text{out}_{i} = \frac{1}{\Gamma(\text{input}_i)}
  \int_{\text{other}_i}^{\infty} t^{\text{input}_i-1} e^{-t} dt

where both \text{input}_i and \text{other}_i are weakly positive
   and at least one is strictly positive. If both are zero or either
   is negative then \text{out}_i=\text{nan}. \Gamma(\cdot) in the
   equation above is the gamma function,
  \Gamma(\text{input}_i) = \int_0^\infty t^{(\text{input}_i-1)}
  e^{-t} dt.

See "torch.special.gammainc()" and "torch.special.gammaln()" for
   related functions.
Supports broadcasting to a common shape and float inputs.
Note:
 The backward pass with respect to "input" is not yet supported.
 Please open an issue on PyTorch's Github to request it.


Parameters:
      * input (Tensor) -- the first non-negative input tensor
  * **other** (*Tensor*) -- the second non-negative input tensor

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a1 = torch.tensor([4.0])
  >>> a2 = torch.tensor([3.0, 4.0, 5.0])
  >>> a = torch.special.gammaincc(a1, a2)
  tensor([0.6472, 0.4335, 0.2650])
  >>> b = torch.special.gammainc(a1, a2) + torch.special.gammaincc(a1, a2)
  tensor([1., 1., 1.])

torch.special.gammaln(input, *, out=None) -> Tensor
Computes the natural logarithm of the absolute value of the gamma
   function on "input".
  \text{out}_{i} = \ln \Gamma(|\text{input}_{i}|)

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a = torch.arange(0.5, 2, 0.5)
  >>> torch.special.gammaln(a)
  tensor([ 0.5724,  0.0000, -0.1208])


tensor([ 0.5724,  0.0000, -0.1208])
torch.special.i0(input, *, out=None) -> Tensor
Computes the zeroth order modified Bessel function of the first
   kind for each element of "input".
  \text{out}_{i} = I_0(\text{input}_{i}) = \sum_{k=0}^{\infty}
  \frac{(\text{input}_{i}^2/4)^k}{(k!)^2}

Parameters:
      input (Tensor) -- the input tensor
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> torch.i0(torch.arange(5, dtype=torch.float32))
  tensor([ 1.0000,  1.2661,  2.2796,  4.8808, 11.3019])

torch.special.i0e(input, *, out=None) -> Tensor
Computes the exponentially scaled zeroth order modified Bessel
   function of the first kind (as defined below) for each element of
   "input".
  \text{out}_{i} = \exp(-|x|) * i0(x) = \exp(-|x|) *
  \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!)^2}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:


Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> torch.special.i0e(torch.arange(5, dtype=torch.float32))
      tensor([1.0000, 0.4658, 0.3085, 0.2430, 0.2070])
torch.special.i1(input, *, out=None) -> Tensor
Computes the first order modified Bessel function of the first kind
   (as defined below) for each element of "input".
  \text{out}_{i} = \frac{(\text{input}_{i})}{2} *
  \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!) *
  (k+1)!}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> torch.special.i1(torch.arange(5, dtype=torch.float32))
      tensor([0.0000, 0.5652, 1.5906, 3.9534, 9.7595])
torch.special.i1e(input, *, out=None) -> Tensor
Computes the exponentially scaled first order modified Bessel
   function of the first kind (as defined below) for each element of
   "input".


"input".
  \text{out}_{i} = \exp(-|x|) * i1(x) =     \exp(-|x|) *
  \frac{(\text{input}_{i})}{2} * \sum_{k=0}^{\infty}
  \frac{(\text{input}_{i}^2/4)^k}{(k!) * (k+1)!}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> torch.special.i1e(torch.arange(5, dtype=torch.float32))
      tensor([0.0000, 0.2079, 0.2153, 0.1968, 0.1788])
torch.special.log1p(input, *, out=None) -> Tensor
Alias for "torch.log1p()".
torch.special.log_ndtr(input, *, out=None) -> Tensor
Computes the log of the area under the standard Gaussian
   probability density function, integrated from minus infinity to
   "input", elementwise.
  \text{log\_ndtr}(x) = \log\left(\frac{1}{\sqrt{2
  \pi}}\int_{-\infty}^{x} e^{-\frac{1}{2}t^2} dt \right)

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.


Example::
      >>> torch.special.log_ndtr(torch.tensor([-3., -2, -1, 0, 1, 2, 3]))
      tensor([-6.6077 -3.7832 -1.841  -0.6931 -0.1728 -0.023  -0.0014])
torch.special.log_softmax(input, dim, *, dtype=None) -> Tensor
Computes softmax followed by a logarithm.
While mathematically equivalent to log(softmax(x)), doing these two
   operations separately is slower and numerically unstable. This
   function is computed as:
  \text{log\_softmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j
  \exp(x_j)} \right)

Parameters:
      * input (Tensor) -- input
  * **dim** (*int*) -- A dimension along which log_softmax will be
    computed.

  * **dtype** ("torch.dtype", optional) -- the desired data type
    of returned tensor. If specified, the input tensor is cast to
    "dtype" before the operation is performed. This is useful for
    preventing data type overflows. Default: None.

Example::
      >>> t = torch.ones(2, 2)


Example::
      >>> t = torch.ones(2, 2)
      >>> torch.special.log_softmax(t, 0)
      tensor([[-0.6931, -0.6931],
              [-0.6931, -0.6931]])
torch.special.logit(input, eps=None, *, out=None) -> Tensor
Returns a new tensor with the logit of the elements of "input".
   "input" is clamped to [eps, 1 - eps] when eps is not None. When eps
   is None and "input" < 0 or "input" > 1, the function will yields
   NaN.
  \begin{align} y_{i} &= \ln(\frac{z_{i}}{1 - z_{i}}) \\ z_{i} &=
  \begin{cases}     x_{i} & \text{if eps is None} \\
  \text{eps} & \text{if } x_{i} < \text{eps} \\     x_{i} &
  \text{if } \text{eps} \leq x_{i} \leq 1 - \text{eps} \\     1 -
  \text{eps} & \text{if } x_{i} > 1 - \text{eps} \end{cases}
  \end{align}

Parameters:
      * input (Tensor) -- the input tensor.
  * **eps** (*float**, **optional*) -- the epsilon for input clamp
    bound. Default: "None"

Keyword Arguments:


Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a = torch.rand(5)
  >>> a
  tensor([0.2796, 0.9331, 0.6486, 0.1523, 0.6516])
  >>> torch.special.logit(a, eps=1e-6)
  tensor([-0.9466,  2.6352,  0.6131, -1.7169,  0.6261])

torch.special.logsumexp(input, dim, keepdim=False, *, out=None)
Alias for "torch.logsumexp()".
torch.special.multigammaln(input, p, *, out=None) -> Tensor
Computes the multivariate log-gamma function with dimension p
   element-wise, given by
  \log(\Gamma_{p}(a)) = C + \displaystyle \sum_{i=1}^{p}
  \log\left(\Gamma\left(a - \frac{i - 1}{2}\right)\right)

where C = \log(\pi) \cdot \frac{p (p - 1)}{4} and \Gamma(-) is the
   Gamma function.
All elements must be greater than \frac{p - 1}{2}, otherwise the
   behavior is undefiend.
Parameters:
      * input (Tensor) -- the tensor to compute the multivariate
        log-gamma function
  * **p** (*int*) -- the number of dimensions


p (int) -- the number of dimensions

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> a = torch.empty(2, 3).uniform_(1, 2)
  >>> a
  tensor([[1.6835, 1.8474, 1.1929],
          [1.0475, 1.7162, 1.4180]])
  >>> torch.special.multigammaln(a, 2)
  tensor([[0.3928, 0.4007, 0.7586],
          [1.0311, 0.3901, 0.5049]])

torch.special.ndtr(input, *, out=None) -> Tensor
Computes the area under the standard Gaussian probability density
   function, integrated from minus infinity to "input", elementwise.
  \text{ndtr}(x) = \frac{1}{\sqrt{2 \pi}}\int_{-\infty}^{x}
  e^{-\frac{1}{2}t^2} dt

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> torch.special.ndtr(torch.tensor([-3., -2, -1, 0, 1, 2, 3]))
      tensor([0.0013, 0.0228, 0.1587, 0.5000, 0.8413, 0.9772, 0.9987])


torch.special.ndtri(input, *, out=None) -> Tensor
Computes the argument, x, for which the area under the Gaussian
   probability density function (integrated from minus infinity to x)
   is equal to "input", elementwise.
  \text{ndtri}(p) = \sqrt{2}\text{erf}^{-1}(2p - 1)

Note:
 Also known as quantile function for Normal Distribution.

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> torch.special.ndtri(torch.tensor([0, 0.25, 0.5, 0.75, 1]))
      tensor([   -inf, -0.6745,  0.0000,  0.6745,     inf])
torch.special.polygamma(n, input, *, out=None) -> Tensor
Computes the n^{th} derivative of the digamma function on "input".
   n \geq 0 is called the order of the polygamma function.
  \psi^{(n)}(x) = \frac{d^{(n)}}{dx^{(n)}} \psi(x)

Note:
 This function is implemented only for nonnegative integers n \geq
 0.

Parameters:


0.
Parameters:
      * n (int) -- the order of the polygamma function
  * **input** (*Tensor*) -- the input tensor.

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> a = torch.tensor([1, 0.5])
      >>> torch.special.polygamma(1, a)
      tensor([1.64493, 4.9348])
      >>> torch.special.polygamma(2, a)
      tensor([ -2.4041, -16.8288])
      >>> torch.special.polygamma(3, a)
      tensor([ 6.4939, 97.4091])
      >>> torch.special.polygamma(4, a)
      tensor([ -24.8863, -771.4742])
torch.special.psi(input, *, out=None) -> Tensor
Alias for "torch.special.digamma()".
torch.special.round(input, *, out=None) -> Tensor
Alias for "torch.round()".
torch.special.scaled_modified_bessel_k0(input, *, out=None) -> Tensor
Scaled modified Bessel function of the second kind of order 0.
Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:


Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
torch.special.scaled_modified_bessel_k1(input, *, out=None) -> Tensor
Scaled modified Bessel function of the second kind of order 1.
Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
torch.special.sinc(input, *, out=None) -> Tensor
Computes the normalized sinc of "input."
  \text{out}_{i} = \begin{cases}   1, & \text{if}\
  \text{input}_{i}=0 \\   \sin(\pi \text{input}_{i}) / (\pi
  \text{input}_{i}), & \text{otherwise} \end{cases}

Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> t = torch.randn(4)
      >>> t
      tensor([ 0.2252, -0.2948,  1.0267, -1.1566])
      >>> torch.special.sinc(t)
      tensor([ 0.9186,  0.8631, -0.0259, -0.1300])


torch.special.softmax(input, dim, *, dtype=None) -> Tensor
Computes the softmax function.
Softmax is defined as:
\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
It is applied to all slices along dim, and will re-scale them so
   that the elements lie in the range [0, 1] and sum to 1.
Parameters:
      * input (Tensor) -- input
  * **dim** (*int*) -- A dimension along which softmax will be
    computed.

  * **dtype** ("torch.dtype", optional) -- the desired data type
    of returned tensor. If specified, the input tensor is cast to
    "dtype" before the operation is performed. This is useful for
    preventing data type overflows. Default: None.

Examples::
      >>> t = torch.ones(2, 2)
      >>> torch.special.softmax(t, 0)
      tensor([[0.5000, 0.5000],
              [0.5000, 0.5000]])
torch.special.spherical_bessel_j0(input, *, out=None) -> Tensor
Spherical Bessel function of the first kind of order 0.
Parameters:


Parameters:
      input (Tensor) -- the input tensor.
Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
torch.special.xlog1py(input, other, *, out=None) -> Tensor
Computes "input * log1p(other)" with the following cases.
  \text{out}_{i} = \begin{cases}     \text{NaN} & \text{if }
  \text{other}_{i} = \text{NaN} \\     0 & \text{if }
  \text{input}_{i} = 0.0 \text{ and } \text{other}_{i} !=
  \text{NaN} \\     \text{input}_{i} *
  \text{log1p}(\text{other}_{i})& \text{otherwise} \end{cases}

Similar to SciPy's scipy.special.xlog1py.
Parameters:
      * input (Number or Tensor) -- Multiplier
  * **other** (*Number** or **Tensor*) -- Argument

Note:
 At least one of "input" or "other" must be a tensor.

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> x = torch.zeros(5,)
  >>> y = torch.tensor([-1, 0, 1, float('inf'), float('nan')])


torch.special.xlog1py(x, y)
      tensor([0., 0., 0., 0., nan])
      >>> x = torch.tensor([1, 2, 3])
      >>> y = torch.tensor([3, 2, 1])
      >>> torch.special.xlog1py(x, y)
      tensor([1.3863, 2.1972, 2.0794])
      >>> torch.special.xlog1py(x, 4)
      tensor([1.6094, 3.2189, 4.8283])
      >>> torch.special.xlog1py(2, y)
      tensor([2.7726, 2.1972, 1.3863])

torch.special.xlogy(input, other, *, out=None) -> Tensor
Computes "input * log(other)" with the following cases.
  \text{out}_{i} = \begin{cases}     \text{NaN} & \text{if }
  \text{other}_{i} = \text{NaN} \\     0 & \text{if }
  \text{input}_{i} = 0.0 \\     \text{input}_{i} *
  \log{(\text{other}_{i})} & \text{otherwise} \end{cases}

Similar to SciPy's scipy.special.xlogy.
Parameters:
      * input (Number or Tensor) -- Multiplier
  * **other** (*Number** or **Tensor*) -- Argument

Note:
 At least one of "input" or "other" must be a tensor.

Keyword Arguments:


Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example:
  >>> x = torch.zeros(5,)
  >>> y = torch.tensor([-1, 0, 1, float('inf'), float('nan')])
  >>> torch.special.xlogy(x, y)
  tensor([0., 0., 0., 0., nan])
  >>> x = torch.tensor([1, 2, 3])
  >>> y = torch.tensor([3, 2, 1])
  >>> torch.special.xlogy(x, y)
  tensor([1.0986, 1.3863, 0.0000])
  >>> torch.special.xlogy(x, 4)
  tensor([1.3863, 2.7726, 4.1589])
  >>> torch.special.xlogy(2, y)
  tensor([2.1972, 1.3863, 0.0000])

torch.special.zeta(input, other, *, out=None) -> Tensor
Computes the Hurwitz zeta function, elementwise.
  \zeta(x, q) = \sum_{k=0}^{\infty} \frac{1}{(k + q)^x}

Parameters:
      * input (Tensor) -- the input tensor corresponding to x.
  * **other** (*Tensor*) -- the input tensor corresponding to *q*.

Note:
 The Riemann zeta function corresponds to the case when *q = 1*

Keyword Arguments:


Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
Example::
      >>> x = torch.tensor([2., 4.])
      >>> torch.special.zeta(x, 1)
      tensor([1.6449, 1.0823])
      >>> torch.special.zeta(x, torch.tensor([1., 2.]))
      tensor([1.6449, 0.0823])
      >>> torch.special.zeta(2, torch.tensor([1., 2.]))
      tensor([1.6449, 0.6449])


torch.utils.bottleneck
torch.utils.bottleneck is a tool that can be used as an initial step
for debugging bottlenecks in your program. It summarizes runs of your
script with the Python profiler and PyTorch's autograd profiler.
Run it on the command line with
python -m torch.utils.bottleneck /path/to/source/script.py [args]
where [args] are any number of arguments to script.py, or run
"python -m torch.utils.bottleneck -h" for more usage instructions.
Warning:
Because your script will be profiled, please ensure that it exits in
  a finite amount of time.
Warning:
Due to the asynchronous nature of CUDA kernels, when running against
  CUDA code, the cProfile output and CPU-mode autograd profilers may
  not show correct timings: the reported CPU time reports the amount
  of time used to launch the kernels but does not include the time the
  kernel spent executing on a GPU unless the operation does a
  synchronize. Ops that do synchronize appear to be extremely


expensive under regular CPU-mode profilers. In these case where
  timings are incorrect, the CUDA-mode autograd profiler may be
  helpful.
Note:
To decide which (CPU-only-mode or CUDA-mode) autograd profiler
  output to look at, you should first check if your script is CPU-
  bound ("CPU total time is much greater than CUDA total time"). If it
  is CPU-bound, looking at the results of the CPU-mode autograd
  profiler will help. If on the other hand your script spends most of
  its time executing on the GPU, then it makes sense to start looking
  for responsible CUDA operators in the output of the CUDA-mode
  autograd profiler.Of course the reality is much more complicated and
  your script might not be in one of those two extremes depending on
  the part of the model you're evaluating. If the profiler outputs
  don't help, you could try looking at the result of
  "torch.autograd.profiler.emit_nvtx()" with "nvprof". However, please
  take into account that the NVTX overhead is very high and often


gives a heavily skewed timeline. Similarly, "IntelÂ® VTuneâ¢ Profiler"
  helps to analyze performance on Intel platforms further with
  "torch.autograd.profiler.emit_itt()".
Warning:
If you are profiling CUDA code, the first profiler that "bottleneck"
  runs (cProfile) will include the CUDA startup time (CUDA buffer
  allocation cost) in its time reporting. This should not matter if
  your bottlenecks result in code much slower than the CUDA startup
  time.
For more complicated uses of the profilers (like in a multi-GPU case),
please see https://docs.python.org/3/library/profile.html or
"torch.autograd.profiler.profile()" for more information.


Frequently Asked Questions
My model reports "cuda runtime error(2): out of memory"
As the error message suggests, you have run out of memory on your GPU.
Since we often deal with large amounts of data in PyTorch, small
mistakes can rapidly cause your program to use up all of your GPU;
fortunately, the fixes in these cases are often simple. Here are a few
common things to check:
Don't accumulate history across your training loop. By default,
computations involving variables that require gradients will keep
history.  This means that you should avoid using such variables in
computations which will live beyond your training loops, e.g., when
tracking statistics. Instead, you should detach the variable or access
its underlying data.
Sometimes, it can be non-obvious when differentiable variables can
occur.  Consider the following training loop (abridged from source):
total_loss = 0
   for i in range(10000):


total_loss = 0
   for i in range(10000):
       optimizer.zero_grad()
       output = model(input)
       loss = criterion(output)
       loss.backward()
       optimizer.step()
       total_loss += loss
Here, "total_loss" is accumulating history across your training loop,
since "loss" is a differentiable variable with autograd history. You
can fix this by writing total_loss += float(loss) instead.
Other instances of this problem: 1.
Don't hold onto tensors and variables you don't need. If you
assign a Tensor or Variable to a local, Python will not deallocate
until the local goes out of scope.  You can free this reference by
using "del x".  Similarly, if you assign a Tensor or Variable to a
member variable of an object, it will not deallocate until the object
goes out of scope.  You will get the best memory usage if you don't
hold onto temporaries you don't need.
The scopes of locals can be larger than you expect.  For example:
for i in range(5):
       intermediate = f(input[i])


intermediate = f(input[i])
       result += g(intermediate)
   output = h(result)
   return output
Here, "intermediate" remains live even while "h" is executing, because
its scope extrudes past the end of the loop.  To free it earlier, you
should "del intermediate" when you are done with it.
Avoid running RNNs on sequences that are too large. The amount of
memory required to backpropagate through an RNN scales linearly with
the length of the RNN input; thus, you will run out of memory if you
try to feed an RNN a sequence that is too long.
The technical term for this phenomenon is backpropagation through
time, and there are plenty of references for how to implement
truncated BPTT, including in the word language model example;
truncation is handled by the "repackage" function as described in this
forum post.
Don't use linear layers that are too large. A linear layer
"nn.Linear(m, n)" uses O(nm) memory: that is to say, the memory


requirements of the weights scales quadratically with the number of
features.  It is very easy to blow through your memory this way (and
remember that you will need at least twice the size of the weights,
since you also need to store the gradients.)
Consider checkpointing. You can trade-off memory for compute by
using checkpoint.
My GPU memory isn't freed properly
PyTorch uses a caching memory allocator to speed up memory
allocations. As a result, the values shown in "nvidia-smi" usually
don't reflect the true memory usage. See Memory management for more
details about GPU memory management.
If your GPU memory isn't freed even after Python quits, it is very
likely that some Python subprocesses are still alive. You may find
them via "ps -elf | grep python" and manually kill them with "kill -9
[pid]".
My out of memory exception handler can't allocate memory


You may have some code that tries to recover from out of memory
errors.
try:
       run_model(batch_size)
   except RuntimeError: # Out of memory
       for _ in range(batch_size):
           run_model(1)
But find that when you do run out of memory, your recovery code can't
allocate either. That's because the python exception object holds a
reference to the stack frame where the error was raised. Which
prevents the original tensor objects from being freed. The solution is
to move you OOM recovery code outside of the "except" clause.
oom = False
   try:
       run_model(batch_size)
   except RuntimeError: # Out of memory
       oom = True
if oom:
       for _ in range(batch_size):
           run_model(1)
My data loader workers return identical random numbers
You are likely using other libraries to generate random numbers in the
dataset and worker subprocesses are started via "fork". See