library directory, and linking "cudart". You can pass additional
   flags to nvcc via "extra_cuda_cflags", just like with
   "extra_cflags" for C++. Various heuristics for finding the CUDA
   install directory are used, which usually work fine. If not,
   setting the "CUDA_HOME" environment variable is the safest option.
Parameters:
      * name -- The name of the extension to build. This MUST be
        the same as the name of the pybind11 module!
  * **sources** (*Union**[**str**, **List**[**str**]**]*) -- A
    list of relative or absolute paths to C++ source files.

  * **extra_cflags** -- optional list of compiler flags to forward
    to the build.

  * **extra_cuda_cflags** -- optional list of compiler flags to
    forward to nvcc when building CUDA sources.

  * **extra_ldflags** -- optional list of linker flags to forward
    to the build.

  * **extra_include_paths** -- optional list of include
    directories to forward to the build.


directories to forward to the build.
  * **build_directory** -- optional path to use as build
    workspace.

  * **verbose** -- If "True", turns on verbose logging of load
    steps.

  * **with_cuda** (*Optional**[**bool**]*) -- Determines whether
    CUDA headers and libraries are added to the build. If set to
    "None" (default), this value is automatically determined based
    on the existence of ".cu" or ".cuh" in "sources". Set it to
    *True`* to force CUDA headers and libraries to be included.

  * **is_python_module** -- If "True" (default), imports the
    produced shared library as a Python module. If "False",
    behavior depends on "is_standalone".

  * **is_standalone** -- If "False" (default) loads the
    constructed extension into the process as a plain dynamic
    library. If "True", build a standalone executable.

Returns:
      Returns the loaded PyTorch extension as a Python module.


If "is_python_module" is "False" and "is_standalone" is "False":
         Returns nothing. (The shared library is loaded into the
         process as a side effect.)
  If "is_standalone" is "True".
     Return the path to the executable. (On Windows,
     TORCH_LIB_PATH is added to the PATH environment variable as a
     side effect.)

Return type:
      If "is_python_module" is "True"
-[ Example ]-

from torch.utils.cpp_extension import load
module = load(
   ...     name='extension',
   ...     sources=['extension.cpp', 'extension_kernel.cu'],
   ...     extra_cflags=['-O2'],
   ...     verbose=True)

torch.utils.cpp_extension.load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False, with_cuda=None, is_python_module=True, with_pytorch_error_handling=True, keep_intermediates=True)


Loads a PyTorch C++ extension just-in-time (JIT) from string
   sources.
This function behaves exactly like "load()", but takes its sources
   as strings rather than filenames. These strings are stored to files
   in the build directory, after which the behavior of "load_inline()"
   is identical to "load()".
See the tests for good examples of using this function.
Sources may omit two required parts of a typical non-inline C++
   extension: the necessary header includes, as well as the (pybind11)
   binding code. More precisely, strings passed to "cpp_sources" are
   first concatenated into a single ".cpp" file. This file is then
   prepended with "#include ".
Furthermore, if the "functions" argument is supplied, bindings will
   be automatically generated for each function specified. "functions"
   can either be a list of function names, or a dictionary mapping
   from function names to docstrings. If a list is given, the name of


each function is used as its docstring.
The sources in "cuda_sources" are concatenated into a separate
   ".cu" file and  prepended with "torch/types.h", "cuda.h" and
   "cuda_runtime.h" includes. The ".cpp" and ".cu" files are compiled
   separately, but ultimately linked into a single library. Note that
   no bindings are generated for functions in "cuda_sources" per  se.
   To bind to a CUDA kernel, you must create a C++ function that calls
   it, and either declare or define this C++ function in one of the
   "cpp_sources" (and include its name in "functions").
See "load()" for a description of arguments omitted below.
Parameters:
      * cpp_sources -- A string, or list of strings, containing
        C++ source code.
  * **cuda_sources** -- A string, or list of strings, containing
    CUDA source code.

  * **functions** -- A list of function names for which to
    generate function bindings. If a dictionary is given, it


should map function names to docstrings (which are otherwise
        just the function names).
  * **with_cuda** -- Determines whether CUDA headers and libraries
    are added to the build. If set to "None" (default), this value
    is automatically determined based on whether "cuda_sources" is
    provided. Set it to "True" to force CUDA headers and libraries
    to be included.

  * **with_pytorch_error_handling** -- Determines whether pytorch
    error and warning macros are handled by pytorch instead of
    pybind. To do this, each function "foo" is called via an
    intermediary "_safe_foo" function. This redirection might
    cause issues in obscure cases of cpp. This flag should be set
    to "False" when this redirect causes issues.

-[ Example ]-

from torch.utils.cpp_extension import load_inline
source = """
   at::Tensor sin_add(at::Tensor x, at::Tensor y) {
     return x.sin() + y.sin();
   }
   """


return x.sin() + y.sin();
   }
   """

module = load_inline(name='inline_extension',
   ...                      cpp_sources=[source],
   ...                      functions=['sin_add'])

Note:
 By default, the Ninja backend uses #CPUS + 2 workers to build the
 extension. This may use up too many resources on some systems.
 One can control the number of workers by setting the *MAX_JOBS*
 environment variable to a non-negative number.

torch.utils.cpp_extension.include_paths(cuda=False)
Get the include paths required to build a C++ or CUDA extension.
Parameters:
      cuda (bool) -- If True, includes CUDA-specific include
      paths.
Returns:
      A list of include path strings.
Return type:
      List[str]
torch.utils.cpp_extension.get_compiler_abi_compatibility_and_version(compiler)
Determine if the given compiler is ABI-compatible with PyTorch
   alongside its version.
Parameters:


alongside its version.
Parameters:
      compiler (str) -- The compiler executable name to check
      (e.g. "g++"). Must be executable in a shell process.
Returns:
      A tuple that contains a boolean that defines if the compiler is
      (likely) ABI-incompatible with PyTorch, followed by a
      TorchVersion string that contains the compiler version
      separated by dots.
Return type:
      Tuple[bool, TorchVersion]
torch.utils.cpp_extension.verify_ninja_availability()
Raises "RuntimeError" if ninja build system is not available on the
   system, does nothing otherwise.
torch.utils.cpp_extension.is_ninja_available()
Returns "True" if the ninja build system is available on the
   system, "False" otherwise.


Installing TorchDynamo
This section describes how to install TorchDynamo. TorchDynamo is
included in the nightly binaries of PyTorch. For more information, see
Getting Started.
Requirements
You must have the following prerequisites to use TorchDynamo:

A Linux or macOS environment

Python 3.8 (recommended). Python 3.7 through 3.10 are supported and
  tested. Make sure to have a development version of Python installed
  locally as well.

GPU/CUDA Requirements
To use GPU back ends, and in particular Triton, make sure that the
CUDA that you have installed locally matches the PyTorch version you
are running.
The following command installs GPU PyTorch + TorchDynamo along with
GPU TorchDynamo dependencies (for CUDA 11.7):
pip3 install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117
CPU requirements
There are no additional requirements for CPU TorchDynamo. CPU


TorchDynamo is included in the nightly versions of PyTorch. To
install, run the following command:
pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu
Install from Local Source
Alternatively, you can build PyTorch from source, which has
TorchDynamo included.
To install GPU TorchDynamo dependencies, run "make triton" in the
PyTorch repo root directory.
Verify Installation
If you built PyTorch from source, then you can run the following
commands (from the PyTorch repo root directory) to check that
TorchDynamo is installed correctly:
cd tools/dynamo
   python verify_dynamo.py
If you do not have the PyTorch source locally, you can alternatively
copy the script ("tools/dynamo/verify_dynamo.py") from the PyTorch
repository and run it locally.
Docker Installation
We also provide all the required dependencies in the PyTorch nightly
binaries which you can download with the following command:


docker pull ghcr.io/pytorch/pytorch-nightly
And for ad hoc experiments just make sure that your container has
access to all your GPUs:
docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash


TorchDynamo Overview
TorchDynamo is a Python-level JIT compiler designed to make
unmodified PyTorch programs faster. TorchDynamo hooks into the frame
evaluation API in CPython (PEP 523) to dynamically modify Python
bytecode right before it is executed. It rewrites Python bytecode in
order to extract sequences of PyTorch operations into an FX Graph
which is then just-in-time compiled with a customizable backend. It
creates this FX Graph through bytecode analysis and is designed to mix
Python execution with compiled backends to get the best of both worlds
â usability and performance.
TorchDynamo makes it easy to experiment with different compiler
backends to make PyTorch code faster with a single line decorator
"torch._dynamo.optimize()"
[image]
TorchInductor is one of the backends supported by TorchDynamo Graph
into Triton for GPUs or C++/OpenMP for CPUs. We have a training
performance dashboard that provides performance comparison for


different training backends. You can read more in the TorchInductor
post on PyTorch dev-discuss.
See also:

TorchDynamo deep-dive video

dev-discuss topics


Guards Overview
From a UX perspective, TorchDynamo is very easy to use. The user
invokes "torchdynamo.optimize" as an annotation:
@torchdynamo.optimize(my_compiler)
   def fn_foo(bar):
Where a complete example looks like this:
from typing import List
   import torch
   import torchdynamo
   def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
       print("my_compiler() called with FX graph:")
       gm.graph.print_tabular()
       return gm.forward  # return a python callable
   @torchdynamo.optimize(my_compiler)
   def toy_example(a, b):
       x = a / (torch.abs(a) + 1)
       if b.sum() < 0:
           b = b * -1
       return x * b
   for _ in range(100):
       toy_example(torch.randn(10), torch.randn(10))
This allows TorchDynamo to capture the interpreted Python frames, grab
any and all relevant information, and speed things up wherever it can.
The speedup comes from a few places, and can be rather dependent on


the backend (my_compiler in the example above) provided, but the one
speedup that is important in this section is caching. Caching
itself is not a direct speedup but a critical enablement that prevents
recompilation. We dig a hole with dynamo, and caching allows us to get
out. It enables us to hold perf neutrality while then enabling
backends - the true source of our speedups.
With even a pass-through no-op backend provided:
def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
       return gm.forward
We can see TorchDynamo speeding up Python execution even on regular
Python, not just PyTorch.
Caching and Guards Overview
TorchDynamo operates through caching transformed (by TorchDynamo) user
bytecode. When TorchDynamo receives a frame for evaluation, it checks
if the objects referenced in the frame have changed in certain
ways, and if not, TorchDynamo reads the previously transformed user


bytecode to evaluate it. In this section, we will focus on how we can
identify whether or not the objects referenced in the frame have
changed. This is a critical piece of functionality in TorchDynamo,
because it drives the entire invalidation lifecycle. This
functionality is called guards.
At a very high level, the flow can be summarized like this:

TorchDynamo receives a Python frame.

It converts the frame (1) passing it through instruction
   translation.

For the objects captured in (2), TorchDynamo creates tracking
   objects that are: * tracked on an output graph, which is an
   internal specialization of a torch.fx.Tracer * guards

TorchDynamo processes the guard objects created in (3), turning
   them into a generated Python function, check_fn, associated with
   a piece of code.

The check_fn is evaluated whenever we encounter this code a
   subsequent time - if a check_fn passes and evaluates to True,
   TorchDynamo identifies the code in the cache and the code


encountered here as same, and can be safely used. If it fails and
   evaluates to False, TorchDynamo identifies the code in the cache
   as not valid, and can be thrown out in favor of a new entry,
   through recompilation or a graph break.
Python Frame Evaluation and PEP 523
The functionality of TorchDynamo is based on PEP 523.
TorchDynamo installs a frame evaluation function on Python by using
_PyInterpreterState_SetEvalFrameFunc. TorchDynamo has a hook where
Python can hand control back to us during evaluation.
The function we have installed is "convert_frame" or
"convert_frame_assert" in the "nopython=True" case, but glossing over
that nuance for now, letâs take a look at "convert_frame_assert", as
"convert_frame" proxies to it.
We can find it on line 20 of convert_frame.py, with a signature as
follows:
def  convert_frame_assert(compiler_fn: Callable, one_graph=True):
This function wraps the entry point of where Python invokes


TorchDynamo with a frame:
def  _convert_frame_assert(frame: types.FrameType, cache_size: int):
Here is what this function does:

Checks if it has seen this "code"(see: f_code here) before and
   exits early if it did.

Checks if the code is an unsupported case.

Checks if the "cache_size" (second arg above) crosses the limit
   defined in the config, "cache_size_limit". If it has, the function
   drops the frame and logs warnings. This helps to avoid constant
   recompilation of a frame as it generally means that the frame is
   hot in an unexpected way and caching it produces needless overhead,
   as it is likely to get evicted the next time it is encountered.

Passes the frame, alongside a function that creates an
   "InstructionTranslator" through bytecode transformation, via
   "transform_code_object". A few crucial things happen under the hood
   here:

New code is produced through "transform_code_object".

An FX tracer named "output" is produced through


"InstructionTranslator".
  This can be a bit confusing, as "InstructionTranslator" is not
  an *fx* tracer, but its stored in a variable named tracer, and
  its output***is***an `fx`tracer.*

The function produces guards and stores them on "output" above.

The function produces "output_instructions" and stores them on
      "output" above.

The function maps the newly produced transformed code to the
      initial code it read off the frame. This mapping is worth
      remembering, we will refer to it much later on below where we
      cover guard failures.

Using the transformed code from 4.1 and the guards from 4.3, the
   function produces a GuardedCode.

Now that we have learned about frame evaluation, letâs review
"InstructionTranslator", and see how it turns the frame we handed it
over into TorchDynamo internal types.
InstructionTranslator
InstructionTranslator does a lot! We wonât cover the details of


everything it does, but most importantly for this document, it
produces a mapping of "symbolic_locals" which maintains a mapping from
the frameâs "f_locals" to TorchDynamo internal Variable objects (more
on these in a moment. "symbolic_locals" is filled via traversing the
frameâs locals:
self.symbolic_locals = collections.OrderedDict(
       (k, VariableBuilder(self, LocalSource(k))(f_locals[k]))
       for k in vars
       if k in f_locals
   )
The important component here  is the invocation of a call into
"VariableBuilder". "VariableBuilder"âs call implementation proxies
into a function called "_wrap", which in turn both constructs
instances of "VariableTracker" and calls "make_guards" on them. More
on that later.
This mapping, in turn, is critical as each Variable has associated
guards, which are then passed to "self.output", the instance of
"OutputGraph", an fx tracer, mentioned in 4.2 of the section above. If
you recall, this "OutputGraph", stored in a variable called "output"


is where our guards are stored before being passed on to become
"GuardedCode"
How does "InstructionTranslator" do this? At the heart of it, there is
a loop that is pumped, which drives a function "step".
"step" is just that - a single processing step, taking exactly one
instruction and doing something with it.
Note:
These are real instructions processed by TorchDynamoâs
  "transform_code_object", and it is pretty cool.
Note:
This section purposely skips the details of dis.get_instructions.
For the example above, here is a snippet of a what a few
"Instruction"'s may look like:
Instruction(opcode=124, opname='LOAD_FAST', arg=0, argval='b', offset=32, starts_line=8, is_jump_target=True, target=None)
   Instruction(opcode=100, opname='LOAD_CONST', arg=3, argval=-1, offset=34, starts_line=None, is_jump_target=False, target=None)
   Instruction(opcode=20, opname='BINARY_MULTIPLY', arg=None, argval=None, offset=36, starts_line=None, is_jump_target=False, target=None)


This is the core functionality of this function. Take a look at the
"opname", and then take a look at this little snippet from inside
"step";
if not hasattr(self, inst.opname):
       unimplemented(f"missing: {inst.opname}")
   getattr(self, inst.opname)(inst)
As we can see, the function checks if the current class, the
"InstructionTranslator" has an attribute set matching the operator
name (for example, "LOAD_CONST"). If it does, the function invokes it,
passing the whole instruction object in. If it does not, the function
drops the frame as unimplemented.
For the "LOAD_CONST" example, we can see that we do indeed support it,
with a relatively straightforward definition:
def  LOAD_CONST(self, inst):
   self.push(ConstantVariable(value=inst.argval))
We can see that this function creates a new instance of the class
"ConstantVariable" , with a value, in our example case, -1, and then
pushes it onto the stack.
There are dozens of such methods - see "symbolic_convert.py" for all


of them. Generally, we implement as many matching methods to Python
bytecode instructions as possible.
Across both the logic downstream of "step" and the logic from invoking
"VariableBuilder" - we now have a lot of "VariableTracker"s and of
course, weâve spoken about creating guards quiet a bit. Letâs dig into
what Variables are, and get a little closer to understanding guards.
Variables
A "ConstantVariable" is an instance of"VariableTracker".
"VariableTracker" represents a tracked Python local or stack value.
When it comes to representing an object inside TorchDynamo, a
"VariableTracker" does exactly what it says - it tracks a given
variable. It is an extremely flexible class, but there are a few
points to keep in mind:

It manages the "guard" relationship around the underlying object
  through:

"make_guard"

"replace_guards"

"add_guard(s)"

"propagate" - "propagate(*vars: List[List["VariableTracker"]])" -


Perhaps the most important of all, in that it combines guards from
    all the provided "VariableTracker" instances passed in. It visits
    the guards and combines the guards from these onto itself.

It acts as a proxy on behalf of the underlying object, implementing
  methods for the rest of TorchDynamo to get information about the
  tracked object:

"call_method"

"call_function"

"python_type"

"as_proxy"

"is/as_python_proxy"

It stores the variable "source" of type "Source", from
  "torchdynamo/source.py". This source type is a relatively self
  contained class that helps us organize and bookkeep where the
  original source came from, and helps provide convenience methods for
  things like getting the name, and importantly for us, producing
  guards.

And this class ("VariableTracker") is built around subclassing,
somewhere between a full Abstract Base Class and fully fleshed out
class - it leaves many methods raising "NotImplementedError" - with


reliance on subclasses. See "torchdynamo/variables/" for all
subclasses to fulfill contracts and custom behaviors.
Knowing what we know now, we can see an example of how an instruction
from "dis", "BUILD_TUPLE":
"BUILD_TUPLE(count)" Creates a tuple consuming count items from the
   stack, and pushes the resulting tuple onto the stack.
In our case, our signature will be a little different due to the way
we create "Instruction" objects, but the gist of it will be the same.
Instead of passing in "count", we pass in an object with a little
extra bookkeeping, and of course, we deal with turning regular old
python objects into TorchDynamo notions:
def BUILD_TUPLE(self, inst):
       items = self.popn(inst.argval)
       options = VariableTracker.propagate(items)
       self.push(TupleVariable(items, **options))
Here is what this code does:

The function reads "argval", which in this case, is analogous to
   "counts" in the pydoc for the equivalent instruction.


The function "popn" the items, in this case, the signature is "def
   popn(self, n: int) -> List[TensorVariable]:" this hints at an
   underlying contract - we are returning "TensorVariables". If we
   take a closer look at "sybmolic_convert.py" and
   "InstructionTranslatorBase"/"InstructionTranslator"we see that the
   only thing pushed onto and popped from our stack are
   "VariableTracker"s.

The function calls "VariableTracker.propagate". This takes the
   guards from every single item popped off the stack in 2, and
   recursively traverses it and combines all the guards into
   "options": "py  return {      "guards": guards,  }"

The function then makes a new instance of a "VariableTracker",
   "TupleVariable"out of the "items" and "options". This then allows
   us to install all the appropriate guards from the "items" that make
   up the new "TupleVariable"

Note:
Where did the first guards come from? Propagation is a good
  technique, but we need something created before it can be


propagated. "VariableBuilder" calls "make_guards" as it creates
  "VariableTracker" instances, from "f_locals". This in turn calls
  into the "source", to have it create guards.
After all this, bytecode translation is done and we are one step
closer to producing "GuardedCode". We now understand how locals become
"VariableTracker"s, how instructions are handled, and where guards are
called on for creation. Before we can go into seeing how code and
guards are combined into a GuardedCode object, we need to dig a little
bit into those "make_guard" and "source.make_guard" calls above. We
can then understand, what was going on when we made guards alongside,
and on, "VariableTracker" instances.
Making Guards
Guards are just Python objects, of the class "Guard". Let's look at
them in more detail.
Looking at the definition of the dataclass (and therefore, ctor
signature), we see that it has a name, a source, and a create
function.
@dataclasses.dataclass
   class Guard:
       name: str


class Guard:
       name: str
       source: GuardSource
       create_fn: Callable
The name should be the name of the variable.
The source here is an enum indicating what kind of source the guard
belongs to.
Note:
Not to be confused with "Source" and the other types in "source.py",
  as stored on "VariableTracker".
"create_fn" provides the main functionality to transition from a
simple dataclass to actually producing valid Python code to be invoked
for knowing whether or not things have changed in between invocations,
and whether we can safely read from the code cache or not.
The most common code paths for getting an instance of a guard are
through "make_guards" on "VariableTracker".
"make_guards"->source.make_guard->return Guard(self.name(),
self.guard_source(), fn)
Or, in a concrete example:
...
   elif istype(value, range):
       guards = self.make_guards(GuardBuilder.EQUALS_MATCH)
       return RangeVariable(value=value, guards=guards)


Since "source" was set at the construction time of this
"VariableTracker", all that was needed here was to provide the "fn",
"GuardBuilder.EQUALS_MATCH" to the "create_fn" field.
This "create_fn" must be a method on "GuardBuilder". The reason for
this becomes apparent in our next step. Once we have all the guards
created for a frame, we move on to "CheckFunctionManager" and
"compile_check_fn".
Before the "convert_frame" function can produce a "GuardedCode", it
needs to run the "CheckFunctionManager", with all the guards, to
produce a "check_fn" which will then, in turn get passed in alongside
the code into "GuardedCode". This is the same "check_fn" that we store
in our cache entry, and the same one we run to know whether or not to
retrieve the code stored alongside. For reference, here is that code:
static CacheEntry create_cache_entry(CacheEntry next,
                                         PyObject guarded_code) {
     CacheEntry e = (CacheEntry *)malloc(sizeof(CacheEntry));


DEBUG_NULL_CHECK(e);
     e->check_fn = PyObject_GetAttrString(guarded_code, "check_fn");
     NULL_CHECK(e->check_fn);
     e->code = (PyCodeObject *)PyObject_GetAttrString(guarded_code, "code");
     NULL_CHECK(e->code);
     e->next = next;
     return e;
   }
We now know how a "check_fn" function is used, and who makes it, and
what it is composed of, but what we do not yet know is how. How does a
list of "Guard" objects become a function we can run later on?
First, we iterate these guards:
for guard in sorted(guards or [], key=Guard.sort_key):
       if not config.guard_nn_modules and guard.is_nn_module():
           continue
       guard.create(local_builder, global_builder)
Calling "guard.create" runs that "create_fn" we set on the "Guard"
class above (donât confuse it with the "check_fn" we are working on
producing, the names are similar, so it can get a little confusing).
In our example above, our "create_fn" is "GuardBuilder.EQUALS_MATCH".


So we are now invoking it, passing in the "self", the guard itself,
in.
The signature is: "def EQUALS_MATCH(self, guard: Guard):"
And internally to that function, we can use the "name" on the guard to
get back our original object, querying it for data and type
information, which in turn gets us to the most important bit:
appending code.
At its simplest, "EQUALS_MATCH" appends just one line of code:
"self.code.append(f"{ref} == {val!r}")". Where "ref" is the name of
the variable, and "val" is the value. It might produce code like this:
y == 2
This is a basic example. But if we append a few other kinds of
"GuardBuilder" functions and then combine them all with "and" in
between each statement (as we do), we might get something like this:
guardedcode.valid and _check_type_id(y, 94367738391392) and y == 2 and ___check_tensors(x)
Here is what this code performs:

A check for ".valid"

A type ID check

A value check

A tensor check


A value check

A tensor check

This becomes the heart of the code our "check_fn", which in turn is
evaluated the next time we encounter this code. It will then
check:

Is this code still valid?

If (1), Does "y" still have a type of "94367738391392"?

If (2), is "y" still 2?

If (3), letâs check on if tensor "x" changed in some specific ways.

If all of these are still true, then we can use the code cached
alongside this "check_fn".
Note:
For a deeper dive for how and where this happens you can read
  "static PyCodeObject lookup(CacheEntry e, PyObject *f_locals) {"
  of "_eval_frame.c".
If not, then, we can move on to recompiling the code anew, and storing
that in the cache alongside this code, and a whole new "check_fn",
again to be checked on yet another subsequent frame.
There are lots of other such functions on "GuardBuilder" which get
coalesced into, at times massive, strings which then get evaluated as
Python code and stored into "check_fn". The example above illustrates


of a simple case. To understand this functionality better, read the
other functions on "GuardBuilder", or better yet, dump the "code"
variable in "compile_check_fn" to see what is getting produced,
especially on larger, real models.
Summary
In this section, we have reviewed:

The role of ".valid" and invalidation around weak references (and
  potentially soon to be NN Moduleinvalidations).

How the C++ side of guard functions ("checktype_id",
  "_check_tensors", etc) operate

What happens when guards fail.

What happens if we produce invalid guard code.

We covered how user provided code wrapped in a TorchDynamo context
goes on to get traced and tracked internally, organized into
"VariableTracker"s "Source"s and subsequently "Guard"s, and how those
"Guards" in turn guide cache entry selection and invalidation when
handing Python code.


DDP Communication Hooks
DDP communication hook is a generic interface to control how to
communicate gradients across workers by overriding the vanilla
allreduce in DistributedDataParallel. A few built-in communication
hooks are provided, and users can easily apply any of these hooks to
optimize communication. Besides, the hook interface can also support
user-defined communication strategies for more advanced use cases.
How to Use a Communication Hook?
To use a communication hook, the user just needs to let the DDP model
register the hook before the training loop as below.
"torch.nn.parallel.DistributedDataParallel.register_comm_hook()"
What Does a Communication Hook Operate On?
A communication hook provides a flexible way to allreduce gradients.
Therefore, it mainly operates on the gradients on each replica before
allreduce, which are bucketized to increase the overlap between


communication and computation. Particularly,
"torch.distributed.GradBucket" represents a bucket of gradient tensors
to be allreduced.
class torch.distributed.GradBucket
This class mainly passes a flattened gradient tensor (returned by
   "buffer()") to DDP communication hook. This tensor can be further
   decomposed into a list of per-parameter tensors within this bucket
   (returned by "get_per_parameter_tensors()") to apply layer-wise
   operations.
torch.distributed.GradBucket.index(self: torch._C._distributed_c10d.GradBucket) -> int
Warning:
 Since the buckets are rebuilt after the first iteration, should
 not rely on the indices at the beginning of training.

Returns:
      The index of a bucket that stores gradients of a few contiguous
      layers. All the gradients are bucketized.
torch.distributed.GradBucket.buffer(self: torch._C._distributed_c10d.GradBucket) -> torch.Tensor
Returns:
      A flattened 1D "torch.Tensor" buffer, which can be further


decomposed into a list of per-parameter tensors within this
      bucket.
torch.distributed.GradBucket.gradients(self: torch._C._distributed_c10d.GradBucket) -> List[torch.Tensor]
Returns:
      A list of "torch.Tensor". Each tensor in the list corresponds to
      a gradient.
torch.distributed.GradBucket.is_last(self: torch._C._distributed_c10d.GradBucket) -> bool
Returns:
      Whether this bucket is the last bucket to allreduce in an
      iteration. This also means that this bucket corresponds to the
      first few layers in the forward pass.
torch.distributed.GradBucket.set_buffer(self: torch._C._distributed_c10d.GradBucket, buffer: torch.Tensor) -> None
Replaces the tensor in the bucket with the input tensor buffer.
torch.distributed.GradBucket.parameters(self: torch._C._distributed_c10d.GradBucket) -> List[torch.Tensor]
Returns:
      A list of "torch.Tensor". Each tensor in the list corresponds to
      a model parameter.
Default Communication Hooks


Default Communication Hooks
Default communication hooks are simple stateless hooks, so the
input state in "register_comm_hook" is either a process group or
"None". The input "bucket" is a "torch.distributed.GradBucket" object.
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook(process_group, bucket)
This DDP communication hook just calls "allreduce" using
   "GradBucket" tensors. Once gradient tensors are aggregated across
   all workers, its "then" callback takes the mean and returns the
   result. If user registers this hook, DDP results is expected to be
   same as the case where no hook was registered. Hence, this won't
   change behavior of DDP and user can use this as a reference or
   modify this hook to log useful information or any other purposes
   while unaffecting DDP behavior.
Example::
      >>> ddp_model.register_comm_hook(process_group, allreduce_hook)
Return type:
      Future[Tensor]


Return type:
      Future[Tensor]
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook(process_group, bucket)
This DDP communication hook implements a simple gradient
   compression approach that casts "GradBucket" tensor to half-
   precision floating-point format ("torch.float16") and then divides
   it by the process group size. It allreduces those "float16"
   gradient tensors. Once compressed gradient tensors are allreduced,
   the chained callback "decompress" casts it back to the input data
   type (such as "float32").
Example::
      >>> ddp_model.register_comm_hook(process_group, fp16_compress_hook)
Return type:
      Future[Tensor]
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook(process_group, bucket)
Warning: This API is experimental, and it requires NCCL version
   later than 2.9.6.
This DDP communication hook implements a simple gradient
   compression approach that casts "GradBucket" tensor to half-


precision Brain floating point format ("torch.bfloat16") and then
   divides it by the process group size. It allreduces those
   "bfloat16" gradient tensors. Once compressed gradient tensors are
   allreduced, the chained callback "decompress" casts it back to the
   input data type (such as "float32").
Example::
      >>> ddp_model.register_comm_hook(process_group, bf16_compress_hook)
Return type:
      Future[Tensor]
Additionally, a communication hook wrapper is provided to support
"fp16_compress_hook()" or "bf16_compress_hook()" as a wrapper, which
can be combined with other communication hooks.
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper(hook)
This wrapper casts the input gradient tensor of a given DDP
   communication hook to half-precision floating point format
   ("torch.float16"), and casts the resulting tensor of the given hook
   back to the input data type, such as "float32".
Therefore, "fp16_compress_hook" is equivalent to


"fp16_compress_wrapper(allreduce_hook)".
Example::
      >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
      >>> ddp_model.register_comm_hook(state, fp16_compress_wrapper(powerSGD_hook))
Return type:
      Callable[[Any, GradBucket], Future[Tensor]]
torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_wrapper(hook)
Warning: This API is experimental, and it requires NCCL version
   later than 2.9.6.
This wrapper casts the input gradient tensor of a given DDP
   communication hook to half-precision Brain floating point format
   https://en.wikipedia.org/wiki/Bfloat16_floating-point_format _
   (``torch.bfloat16), and casts the resulting tensor of the given
   hook back to the input data type, such as "float32".
Therefore, "bf16_compress_hook" is equivalent to
   "bf16_compress_wrapper(allreduce_hook)".
Example::


Example::
      >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1, start_powerSGD_iter=10)
      >>> ddp_model.register_comm_hook(state, bf16_compress_wrapper(powerSGD_hook))
Return type:
      Callable[[Any, GradBucket], Future[Tensor]]
PowerSGD Communication Hook
PowerSGD (Vogels et al., NeurIPS 2019) is a gradient compression
algorithm, which can provide very high compression rates and
accelerate bandwidth-bound distributed training. This algorithm needs
to maintain both some hyperparameters and the internal state.
Therefore, PowerSGD communication hook is a stateful hook, and the
user needs to provide a state object defined as below.
PowerSGD State


PowerSGD State
class torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState(process_group, matrix_approximation_rank=1, start_powerSGD_iter=1000, min_compression_rate=2, use_error_feedback=True, warm_start=True, orthogonalization_epsilon=0, random_seed=0, compression_stats_logging_frequency=10000, batch_tensors_with_same_shape=False)
Stores both the algorithm's hyperparameters and the internal state
   for all the gradients during the training. Particularly,
   "matrix_approximation_rank" and "start_powerSGD_iter" are the main
   hyperparameters that should be tuned by the user. For performance,
   we suggest to keep binary hyperparameters "use_error_feedback" and
   "warm_start" on.

"matrix_approximation_rank" controls the size of compressed low-
      rank tensors, which determines the compression rate. The lower
      the rank, the stronger the compression. 1.1. If "matrix_approximation_rank" is too low, the full


model quality will need more training steps to reach or will
         never reach and yield loss in accuracy.
     1.2. The increase of "matrix_approximation_rank" can
     substantially increase the computation costs of the
     compression, and the accuracy may not be further improved
     beyond a certain "matrix_approximation_rank" threshold.

To tune "matrix_approximation_rank", we suggest to start from 1 and
   increase by factors of 2 (like an exponential grid search, 1, 2, 4,
   ...), until a satisfactory accuracy is reached. Typically only a
   small value 1-4 is used. For some NLP tasks (as shown in Appendix D
   of the original paper), this value has been increased to 32.

"start_powerSGD_iter" defers PowerSGD compression until step
      "start_powerSGD_iter", and vanilla allreduce runs prior to step
      "start_powerSGD_iter". This hybrid scheme of **vanilla allreduce
PowerSGD** can effectively improve the accuracy, even a


relatively small "matrix_approximation_rank" is used. This is
      because that, the beginning of training phase is usually very
      sensitive to inaccurate gradients, and compressing gradients too
      early may make the training quickly take a suboptimal
      trajectory, which can result in an irrecoverable impact on the
      accuracy.
To tune "start_powerSGD_iter", we suggest to start with 10% of
   total training steps, and increase it until a satisfactory accuracy
   is reached. If there is a warm-up stage in the training,
   "start_powerSGD_iter" typically should be no less than the number
   of warm-up steps.

"min_compression_rate" is the minimum compression rate required
      when a layer is compressed. Due to the computation overheads
      incurred by the compression, a tensor is worth compressing only
      if there can be sufficient saving in bandwidth, where "(num_rows
num_cols) * matrix_approximation_rank * min_compression_rate <


num_rows * num_cols". If the specified compression rate
      threshold cannot be satisfied, the tensor will be directly
      allreduced without compression.
Compression statistics are logged every
   "compression_stats_logging_frequency" iterations once PowerSGD
   compression starts.

"orthogonalization_epsilon" can be a very small value (e.g.,
      1e-8) added to every normalized matrix column in
      orthogonalization step, to prevent div-by-zero error if any
      column has all 0s. If this can already be prevented (e.g., by
      batch normalization), an epsilon of 0 is recommended for
      accuracy.

"batch_tensors_with_same_shape" controls whether to compress and
      decompress tensors with same shape in a batched operation to
      achieve higher parallelism. Note that you should also increase
      the bucket size (i.e., "bucket_cap_mb" arg in DDP constructor)
      to make more same-shaped tensors appear in the same bucket,


however this may reduce the overlap between computation and
      communication, and increase the memory footprint due to stacking
      the tensors of the same shape. Set to "True" if the compression
      / decompression computation is a bottleneck.
Warning:
 If error feedback or warm-up is enabled, the minimum value of
 "start_powerSGD_iter" allowed in DDP is 2. This is because there
 is another internal optimization that rebuilds buckets at
 iteration 1 in DDP, and this can conflict with any tensor
 memorized before the rebuild process.

PowerSGD Hooks
Warning:
PowerSGD typically requires extra memory of the same size as the
  model's gradients to enable error feedback, which can compensate for
  biased compressed communication and improve accuracy.
Warning:
PowerSGD hooks may conflict with Apex automatic mixed precision
  package. Please use PyTorch native automatic mixed precision package
  instead.


instead.
torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook(state, bucket)
This DDP communication hook implements PowerSGD gradient
   compression algorithm described in the paper. Once gradient tensors
   are aggregated across all workers, this hook applies compression as
   follows:

Views the input flattened 1D gradient tensor as a list of per-
      parameter tensors, and divides all the tensors into two groups:
 1.1 The tensors that should be compressed before allreduce,
 because the compression can give enough saving in bandwidth.

 1.2 Rest of the tensors will be directly allreduced without
 compression, including all the vector tensors (for biases).

Handles uncompressed tensors:
 2.1. Allocate contiguous memory for those uncompressed
 tensors, and allreduces all the uncompressed tensors as a
 batch, without compression;

 2.2. Copies the individual uncompressed tensors from the


contiguous memory back to the input tensor.

Handles the tensors that should be compressed by PowerSGD
      compression: 3.1. For each tensor M, creates two low-rank tensors P and Q
 for decomposing M, such that M = PQ^T, where Q is initialized
 from a standard normal distribution and orthogonalized;

 3.2. Computes each P in Ps, which is equal to MQ;

 3.3. Allreduces Ps as a batch;

 3.4. Orthogonalizes each P in Ps;

 3.5. Computes each Q in Qs, which is approximately equal to
 M^TP;

 3.6. Allreduces Qs as a batch;

 3.7. Computes each M among all the compressed tensors, which
 is approximately equal to PQ^T.

Note that this communication hook enforces vanilla allreduce for
   the first "state.start_powerSGD_iter" iterations. This not only
   gives the user more control over the tradeoff between speedup and
   accuracy, but also helps abstract away some complexity of the


internal optimization of DDP for future communication hook
   developers.
Parameters:
      * state (PowerSGDState) -- State information to configure
        the compression rate and support error feedback, warm start,
        etc. To tune the compression configs, mainly need to tune
        "matrix_approximation_rank", "start_powerSGD_iter" and
        "min_compression_rate".
  * **bucket** (*dist.GradBucket*) -- Bucket that stores a 1D
    flattened gradient tensor that batches multiple per-variable
    tensors. Note that since DDP comm hook only supports single
    process single device mode, only exactly one tensor is stored
    in this bucket.

Returns:
      Future handler of the communication, which updates the gradients
      in place.
Return type:
      Future[Tensor]
Example::
      >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1,
                                start_powerSGD_iter=10, min_compression_rate=0.5)


ddp_model.register_comm_hook(state, powerSGD_hook)

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.batched_powerSGD_hook(state, bucket)
This DDP communication hook implements a simplified PowerSGD
   gradient compression algorithm described in the paper. This variant
   does not compress the gradients layer by layer, but instead
   compresses the flattened input tensor that batches all the
   gradients. Therefore, it is faster than "powerSGD_hook()", but
   usually results in a much lower accuracy, unless
   "matrix_approximation_rank" is 1.
Warning:
 Increasing "matrix_approximation_rank" here may not necessarily
 increase the accuracy, because batching per-parameter tensors
 without column/row alignment can destroy low-rank structure.
 Therefore, the user should always consider "powerSGD_hook()"
 first, and only consider this variant when a satisfactory
 accuracy can be achieved when "matrix_approximation_rank" is 1.


Once gradient tensors are aggregated across all workers, this hook
   applies compression as follows:

Views the input flattened 1D gradient tensor as a square-shaped
      tensor M with 0 paddings;

Creates two low-rank tensors P and Q for decomposing M, such
      that M = PQ^T, where Q is initialized from a standard normal
      distribution and orthogonalized;

Computes P, which is equal to MQ;

Allreduces P;

Orthogonalizes P;

Computes Q, which is approximately equal to M^TP;

Allreduces Q;

Computes M, which is approximately equal to PQ^T.

Truncates the input tensor to the original length.

Note that this communication hook enforces vanilla allreduce for
   the first "state.start_powerSGD_iter" iterations. This not only
   gives the user more control over the tradeoff between speedup and
   accuracy, but also helps abstract away some complexity of the
   internal optimization of DDP for future communication hook
   developers.


developers.
Parameters:
      * state (PowerSGDState) -- State information to configure
        the compression rate and support error feedback, warm start,
        etc. To tune the compression configs, mainly need to tune
        "matrix_approximation_rank" and "start_powerSGD_iter".
  * **bucket** (*dist.GradBucket*) -- Bucket that stores a 1D
    flattened gradient tensor that batches multiple per-variable
    tensors. Note that since DDP comm hook only supports single
    process single device mode, only exactly one tensor is stored
    in this bucket.

Returns:
      Future handler of the communication, which updates the gradients
      in place.
Return type:
      Future[Tensor]
Example::
      >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1)
      >>> ddp_model.register_comm_hook(state, batched_powerSGD_hook)
Debugging Communication Hooks


=============================
As the name implies, debugging communication hooks are only used
for debugging and performance optimization purpose.
Warning:
Debugging communication hooks do not necessarily output the correct
  results.
torch.distributed.algorithms.ddp_comm_hooks.debugging_hooks.noop_hook(_, bucket)
This DDP communication hook returns a future that wraps the input,
   so it is a noop that does not incur any communication overheads.
This hook should only be used for headroom analysis of
   allreduce optimization, instead of the normal gradient
   synchronization. For example, if only less than 10% speedup of
   training time can be observed after this hook is registered, it
   usually implies that allreduce is not a performance bottleneck for
   this case. Such instrumentation can be particularly useful if GPU
   traces cannot be easily retrieved or the trace analysis is
   complicated some factors such as the overlap between allreduce and


computation or the desynchronization across ranks.
Example::
      >>> ddp_model.register_comm_hook(None, noop_hook)
Return type:
      Future[Tensor]
Checkpointing of Communication Hooks
A stateful communication hook can be saved as a part of model
checkpointing to enable trainer restarts. To make a hook serializable,
"setstate" and "getstate" should be defined.
Warning:
"getstate" should exclude non-serializable attributes from a
  returned dictionary.
Warning:
"setstate" should properly initialize non-serializable
  attributes, excluded from a provided "state".
"PowerSGDState" has "setstate" and "getstate" implemented and
can be used as a reference.
class torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState
getstate()
  Returns a "Dict[str, Any]" which will be pickled and saved.
  "process_group" is not serializable and excluded from a returned
  state.

setstate(state)


state.
setstate(state)
  Takes a provided "state" and retrieves "PowerSGDState".
  "process_group" is set to default.

Here is a simple, end-to-end example of saving and reloading PowerSGD
state and hook.
import os
   import sys
   import tempfile
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   import torch.optim as optim
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook as powerSGD
class SimpleModel(nn.Module):
       def init(self):
           super(SimpleModel, self).init()
           self.fc1 = nn.Linear(24,24)
           self.relu = nn.ReLU()
           self.fc2 = nn.Linear(24,12)
   def forward(self, x):
       return self.fc2(self.relu(self.fc1(x)))

def setup(rank, world_size):
       os.environ['MASTER_ADDR'] = 'localhost'
       os.environ['MASTER_PORT'] = '12355'
   # initialize the process group
   dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():


def cleanup():
       dist.destroy_process_group()
def run_demo(demo_fn, world_size):
       mp.spawn(
           demo_fn,
           args=(world_size,),
           nprocs=world_size,
           join=True)
def demo_serialization(rank, world_size):
       setup(rank, world_size)
   CHECKPOINT = tempfile.gettempdir() + "/checkpoint.pt"

   model = SimpleModel().to(rank)
   ddp_model = DistributedDataParallel(model, device_ids=[rank])

   powersgd_hook = powerSGD.powerSGD_hook
   powersgd_state = powerSGD.PowerSGDState(process_group=None)

   optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
   ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

   state = {
       'state_dict': ddp_model.state_dict(),
       'comm_hook': hook,
       'comm_hook_state': hook_state}

   if rank == 0:
       torch.save(state, CHECKPOINT)

   dist.barrier()
   map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}


checkpoint = torch.load(CHECKPOINT, map_location=map_location)
   ddp_model.load_state_dict(checkpoint['state_dict'])
   powersgd_hook = checkpoint['comm_hook']
   powersgd_state = checkpoint['comm_hook_state']

   ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

   if rank == 0:
       os.remove(CHECKPOINT)

   cleanup()

if name == "main":
       n_gpus = torch.cuda.device_count()
       assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
       world_size = n_gpus
       run_demo(demo_serialization, world_size)
Acknowledgements
Many thanks to PowerSGD paper author Thijs Vogels for the code
review on PowerSGD communication hook, as well as the comparison
experiments, which show that the performance of PowerSGD communication
hook is on par with the implementation in the original paper.


Pipeline Parallelism
Pipeline parallelism was original introduced in the Gpipe  paper and
is an efficient technique to train large models on multiple GPUs.
Warning:
Pipeline Parallelism is experimental and subject to change.
Model Parallelism using multiple GPUs
Typically for large models which don't fit on a single GPU, model
parallelism is employed where certain parts of the model are placed on
different GPUs. Although, if this is done naively for sequential
models, the training process suffers from GPU under utilization since
only one GPU is active at one time as shown in the figure below:
[image]The figure represents a model with 4 layers placed on 4
   different GPUs (vertical axis). The horizontal axis represents
   training this model through time demonstrating that only 1 GPU is
   utilized at a time (image source).
Pipelined Execution
To alleviate this problem, pipeline parallelism splits the input


minibatch into multiple microbatches and pipelines the execution of
these microbatches across multiple GPUs. This is outlined in the
figure below:
[image]The figure represents a model with 4 layers placed on 4
   different GPUs (vertical axis). The horizontal axis represents
   training this model through time demonstrating that the GPUs are
   utilized much more efficiently. However, there still exists a
   bubble (as demonstrated in the figure) where certain GPUs are not
   utilized. (image source).
Pipe APIs in PyTorch
class torch.distributed.pipeline.sync.Pipe(module, chunks=1, checkpoint='except_last', deferred_batch_norm=False)
Wraps an arbitrary "nn.Sequential" module to train on using
   synchronous pipeline parallelism. If the module requires lots of
   memory and doesn't fit on a single GPU, pipeline parallelism is a
   useful technique to employ for training.
The implementation is based on the torchgpipe paper.


Pipe combines pipeline parallelism with checkpointing to reduce
   peak memory required to train while minimizing device under-
   utilization.
You should place all the modules on the appropriate devices and
   wrap them into an "nn.Sequential" module defining the desired order
   of execution. If a module does not contain any parameters/buffers,
   it is assumed this module should be executed on CPU and appropriate
   input tensors to the module are moved to CPU before execution. This
   behavior can be overridden by the "WithDevice" wrapper which can be
   used to explicitly specify which device a module should run on.
Parameters:
      * module ("nn.Sequential") -- sequential module to be
        parallelized using pipelining. Each module in the sequence has
        to have all of its parameters on a single device. Each module
        in the sequence has to either be an nn.Module or
        "nn.Sequential" (to combine multiple sequential modules on a
        single device)


single device)
  * **chunks** (*int*) -- number of micro-batches (default: "1")

  * **checkpoint** (*str*) -- when to enable checkpointing, one of
    "'always'", "'except_last'", or "'never'" (default:
    "'except_last'"). "'never'" disables checkpointing completely,
    "'except_last'" enables checkpointing for all micro-batches
    except the last one and "'always'" enables checkpointing for
    all micro-batches.

  * **deferred_batch_norm** (*bool*) -- whether to use deferred
    "BatchNorm" moving statistics (default: "False"). If set to
    "True", we track statistics across multiple micro-batches to
    update the running statistics per mini-batch.

Raises:
      * TypeError -- the module is not a "nn.Sequential".
  * **ValueError** -- invalid arguments

Example::
      Pipeline of two FC layers across GPUs 0 and 1.
  >>> # Need to initialize RPC framework first.
  >>> os.environ['MASTER_ADDR'] = 'localhost'


os.environ['MASTER_ADDR'] = 'localhost'
      >>> os.environ['MASTER_PORT'] = '29500'
      >>> torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1)
      >>>
      >>> # Build pipe.
      >>> fc1 = nn.Linear(16, 8).cuda(0)
      >>> fc2 = nn.Linear(8, 4).cuda(1)
      >>> model = nn.Sequential(fc1, fc2)
      >>> model = Pipe(model, chunks=8)
      >>> input = torch.rand(16, 16).cuda(0)
      >>> output_rref = model(input)

Note:
 You can wrap a "Pipe" model with
 "torch.nn.parallel.DistributedDataParallel" only when the
 checkpoint parameter of "Pipe" is "'never'".

Note:
 "Pipe" only supports intra-node pipelining currently, but will be
 expanded to support inter-node pipelining in the future. The
 forward function returns an "RRef" to allow for inter-node
 pipelining in the future, where the output might be on a remote
 host. For intra-node pipelinining you can use "local_value()" to
 retrieve the output locally.

Warning:


retrieve the output locally.
Warning:
 "Pipe" is experimental and subject to change.

forward(*inputs)
  Processes a single input mini-batch through the pipe and returns
  an "RRef" pointing to the output. "Pipe" is a fairly transparent
  module wrapper. It doesn't modify the input and output signature
  of the underlying module. But there's type restriction. Input
  and output have to contain at least one tensor. This restriction
  is applied at partition boundaries too.

  The sequence of inputs are fed into the first stage of the
  pipeline as "*inputs". As a result the positional args for this
  function should match the positional args for the first stage of
  the pipeline. The same condition applies for output of one stage
  of the pipeline which is the input for the next stage.

  The input tensor is split into multiple micro-batches based on
  the "chunks" parameter used to initialize "Pipe". The batch size


is assumed to be the first dimension of the tensor and if the
      batch size is less than "chunks", the number of micro-batches is
      equal to the batch size.
  Only tensors are split into multiple micro-batches, non-Tensor
  inputs are just replicated as-is in each micro-batch. For non-
  Tensor outputs in the last stage of the pipeline, they are
  aggregated as a "List" and returned the user. For example, if
  you have 2 micro-batches returning the integer 5, the user would
  receive the consolidated output of *[5, 5]*

  All the input tensors need to be on the same device as the first
  partition of the pipeline.

  If a tensor is wrapped with the "NoChunk" wrapper, the tensor is
  not split across micro-batches and is replicated as-is similar
  to non-tensors.

  Parameters:
     **inputs** -- input mini-batch

  Returns:
     "RRef" to the output of the mini-batch

  Raises:


Raises:
         TypeError -- input doesn't contain at least one tensor
  Return type:
     *RRef*

Skip connections
Certain models like ResNeXt are not completely sequential and have
skip connections between layers. Naively implementing as part of
pipeline parallelism would imply that we need to copy outputs for
certain layers through multiple GPUs till we eventually reach the GPU
where the layer for the skip connection resides. To avoid this copy
overhead, we provide APIs below to stash and pop Tensors in different
layers of the model.
torch.distributed.pipeline.sync.skip.skippable.skippable(stash=(), pop=())
The decorator to define a "nn.Module" with skip connections.
   Decorated modules are called "skippable". This functionality works
   perfectly fine even when the module is not wrapped by "Pipe".
Each skip tensor is managed by its name. Before manipulating skip
   tensors, a skippable module must statically declare the names for


skip tensors by stash and/or pop parameters. Skip tensors with
   pre-declared name can be stashed by "yield stash(name, tensor)" or
   popped by "tensor = yield pop(name)".
Here is an example with three layers. A skip tensor named "1to3" is
   stashed and popped at the first and last layer, respectively:
  @skippable(stash=['1to3'])
  class Layer1(nn.Module):
      def forward(self, input):
          yield stash('1to3', input)
          return f1(input)

  class Layer2(nn.Module):
      def forward(self, input):
          return f2(input)

  @skippable(pop=['1to3'])
  class Layer3(nn.Module):
      def forward(self, input):
          skip_1to3 = yield pop('1to3')
          return f3(input) + skip_1to3

  model = nn.Sequential(Layer1(), Layer2(), Layer3())

One skippable module can stash or pop multiple skip tensors:
  @skippable(stash=['alice', 'bob'], pop=['carol'])
  class StashStashPop(nn.Module):


class StashStashPop(nn.Module):
          def forward(self, input):
              yield stash('alice', f_alice(input))
              yield stash('bob', f_bob(input))
              carol = yield pop('carol')
              return input + carol
Every skip tensor must be associated with exactly one pair of
   stash and pop. "Pipe" checks this restriction automatically
   when wrapping a module. You can also check the restriction by
   "verify_skippables()" without "Pipe".
Return type:
      Callable[[Type[Module]], Type[Skippable]]
class torch.distributed.pipeline.sync.skip.skippable.stash(name, tensor)
The command to stash a skip tensor.
  def forward(self, input):
      yield stash('name', input)
      return f(input)

Parameters:
      * name (str) -- name of skip tensor
  * **input** (*torch.Tensor** or **None*) -- tensor to pass to
    the skip connection

class torch.distributed.pipeline.sync.skip.skippable.pop(name)


The command to pop a skip tensor.
  def forward(self, input):
      skip = yield pop('name')
      return f(input) + skip

Parameters:
      name (str) -- name of skip tensor
Returns:
      the skip tensor previously stashed by another layer under the
      same name
Return type:
      None
torch.distributed.pipeline.sync.skip.skippable.verify_skippables(module)
Verifies if the underlying skippable modules satisfy integrity.
Every skip tensor must have only one pair of stash and pop. If
   there are one or more unmatched pairs, it will raise "TypeError"
   with the detailed messages.
Here are a few failure cases. "verify_skippables()" will report
   failure for these cases:
  # Layer1 stashes "1to3".
  # Layer3 pops "1to3".

  nn.Sequential(Layer1(), Layer2())
  #               âââââ ?

  nn.Sequential(Layer2(), Layer3())
  #                   ? âââââ

  nn.Sequential(Layer1(), Layer2(), Layer3(), Layer3())


âââââââââââââââââââââ       ^^^^^^
  nn.Sequential(Layer1(), Layer1(), Layer2(), Layer3())
  #             ^^^^^^      âââââââââââââââââââââ

To use the same name for multiple skip tensors, they must be
   isolated by different namespaces. See "isolate()".
Raises:
      TypeError -- one or more pairs of stash and pop are not
      matched.
Tutorials
The following tutorials give a good overview of how to use the "Pipe"
API to train your models with the rest of the components that PyTorch
provides:

Training Transformer models using Pipeline Parallelism

Training Transformer models using Distributed Data Parallel and
  Pipeline Parallelism

Acknowledgements
The implementation for pipeline parallelism is based on fairscale's
pipe implementation and torchgpipe. We would like to thank both teams


for their contributions and guidance towards bringing pipeline
parallelism into PyTorch.


Distributed Checkpoint


torch.backends
torch.backends controls the behavior of various backends that
PyTorch supports.
These backends include:

"torch.backends.cuda"

"torch.backends.cudnn"

"torch.backends.mps"

"torch.backends.mkl"

"torch.backends.mkldnn"

"torch.backends.openmp"

"torch.backends.opt_einsum"

"torch.backends.xeon"

torch.backends.cuda
torch.backends.cuda.is_built()
Returns whether PyTorch is built with CUDA support.  Note that this
   doesn't necessarily mean CUDA is available; just that if this
   PyTorch binary were run a machine with working CUDA drivers and
   devices, we would be able to use it.
torch.backends.cuda.matmul.allow_tf32
A "bool" that controls whether TensorFloat-32 tensor cores may be
   used in matrix multiplications on Ampere or newer GPUs. See
   TensorFloat-32(TF32) on Ampere devices.
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction
A "bool" that controls whether reduced precision reductions (e.g.,


with fp16 accumulation type) are allowed with fp16 GEMMs.
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction
A "bool" that controls whether reduced precision reductions are
   allowed with bf16 GEMMs.
torch.backends.cuda.cufft_plan_cache
"cufft_plan_cache" caches the cuFFT plans
size
  A readonly "int" that shows the number of plans currently in the
  cuFFT plan cache.

torch.backends.cuda.max_size
  A "int" that controls cache capacity of cuFFT plan.

torch.backends.cuda.clear()
  Clears the cuFFT plan cache.

torch.backends.cuda.preferred_linalg_library(backend=None)
Warning:
 This flag is experimental and subject to change.

When PyTorch runs a CUDA linear algebra operation it often uses the
   cuSOLVER or MAGMA libraries, and if both are available it decides
   which to use with a heuristic. This flag (a "str") allows
   overriding those heuristics.

If "cusolver" is set then cuSOLVER will be used wherever
     possible.


possible.

If "magma" is set then MAGMA will be used wherever possible.

If "default" (the default) is set then heuristics will be used
     to pick between cuSOLVER and MAGMA if both are available.

When no input is given, this function returns the currently
     preferred library.

Note: When a library is preferred other libraries may still be used
   if the preferred library doesn't implement the operation(s) called.
   This flag may achieve better performance if PyTorch's heuristic
   library selection is incorrect for your application's inputs.
Currently supported linalg operators:

"torch.linalg.inv()"

"torch.linalg.inv_ex()"

"torch.linalg.cholesky()"

"torch.linalg.cholesky_ex()"

"torch.cholesky_solve()"

"torch.cholesky_inverse()"

"torch.linalg.lu_factor()"

"torch.linalg.lu()"

"torch.linalg.lu_solve()"

"torch.linalg.qr()"

"torch.linalg.eigh()"

"torch.linalg.eighvals()"

"torch.linalg.svd()"


"torch.linalg.svd()"

"torch.linalg.svdvals()"

Return type:
      _LinalgBackend
class torch.backends.cuda.SDPBackend(value)
Enum class for the scaled dot product attention backends.
Warning:
 This flag is experimental and subject to change.'

This class needs to stay inline with the enum defined in:
   pytorch/aten/src/ATen/native/transformers/sdp_utils_cpp.h
torch.backends.cuda.flash_sdp_enabled()
Warning:
 This flag is experimental and subject to change.

Returns whether flash sdp is enabled or not.
torch.backends.cuda.enable_mem_efficient_sdp(enabled)
Warning:
 This flag is experimental and subject to change.

Enables or disables memory efficient sdp.
torch.backends.cuda.mem_efficient_sdp_enabled()
Warning:
 This flag is experimental and subject to change.

Returns whether memory efficient sdp is enabled or not.
torch.backends.cuda.enable_flash_sdp(enabled)
Warning:
 This flag is experimental and subject to change.


Enables or disables flash sdp.
torch.backends.cuda.math_sdp_enabled()
Warning:
 This flag is experimental and subject to change.

Returns whether math sdp is enabled or not.
torch.backends.cuda.enable_math_sdp(enabled)
Warning:
 This flag is experimental and subject to change.

Enables or disables math sdp.
torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True)
Warning:
 This flag is experimental and subject to change.

This context manager can be used to temporarily enable or disable
   flash/memory efficient sdp and math sdp. Upon exiting the context
   manager, the previous state of the flags will be restored.
torch.backends.cudnn
torch.backends.cudnn.version()
Returns the version of cuDNN
torch.backends.cudnn.is_available()
Returns a bool indicating if CUDNN is currently available.
torch.backends.cudnn.enabled
A "bool" that controls whether cuDNN is enabled.


torch.backends.cudnn.allow_tf32
A "bool" that controls where TensorFloat-32 tensor cores may be
   used in cuDNN convolutions on Ampere or newer GPUs. See
   TensorFloat-32(TF32) on Ampere devices.
torch.backends.cudnn.deterministic
A "bool" that, if True, causes cuDNN to only use deterministic
   convolution algorithms. See also
   "torch.are_deterministic_algorithms_enabled()" and
   "torch.use_deterministic_algorithms()".
torch.backends.cudnn.benchmark
A "bool" that, if True, causes cuDNN to benchmark multiple
   convolution algorithms and select the fastest.
torch.backends.cudnn.benchmark_limit
A "int" that specifies the maximum number of cuDNN convolution
   algorithms to try when torch.backends.cudnn.benchmark is True.
   Set benchmark_limit to zero to try every available algorithm.
   Note that this setting only affects convolutions dispatched via the
   cuDNN v8 API.
torch.backends.mps
torch.backends.mps.is_available()


torch.backends.mps.is_available()
Returns a bool indicating if MPS is currently available.
Return type:
      bool
torch.backends.mps.is_built()
Returns whether PyTorch is built with MPS support. Note that this
   doesn't necessarily mean MPS is available; just that if this
   PyTorch binary were run a machine with working MPS drivers and
   devices, we would be able to use it.
Return type:
      bool
torch.backends.mkl
torch.backends.mkl.is_available()
Returns whether PyTorch is built with MKL support.
class torch.backends.mkl.verbose(enable)
On-demand oneMKL verbosing functionality To make it easier to debug
   performance issues, oneMKL can dump verbose messages containing
   execution information like duration while executing the kernel. The
   verbosing functionality can be invoked via an environment variable
   named MKL_VERBOSE. However, this methodology dumps messages in
   all steps. Those are a large amount of verbose messages. Moreover,


for investigating the performance issues, generally taking verbose
   messages for one single iteration is enough. This on-demand
   verbosing functionality makes it possible to control scope for
   verbose message dumping. In the following example, verbose messages
   will be dumped out for the second inference only.
  import torch
  model(data)
  with torch.backends.mkl.verbose(torch.backends.mkl.VERBOSE_ON):
      model(data)

Parameters:
      level -- Verbose level - "VERBOSE_OFF": Disable verbosing -
      "VERBOSE_ON":  Enable verbosing
torch.backends.mkldnn
torch.backends.mkldnn.is_available()
Returns whether PyTorch is built with MKL-DNN support.
class torch.backends.mkldnn.verbose(level)
On-demand oneDNN (former MKL-DNN) verbosing functionality To make
   it easier to debug performance issues, oneDNN can dump verbose
   messages containing information like kernel size, input data size


and execution duration while executing the kernel. The verbosing
   functionality can be invoked via an environment variable named
   DNNL_VERBOSE. However, this methodology dumps messages in all
   steps. Those are a large amount of verbose messages. Moreover, for
   investigating the performance issues, generally taking verbose
   messages for one single iteration is enough. This on-demand
   verbosing functionality makes it possible to control scope for
   verbose message dumping. In the following example, verbose messages
   will be dumped out for the second inference only.
  import torch
  model(data)
  with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
      model(data)

Parameters:
      level -- Verbose level - "VERBOSE_OFF": Disable verbosing -
      "VERBOSE_ON":  Enable verbosing - "VERBOSE_ON_CREATION": Enable
      verbosing, including oneDNN kernel creation
torch.backends.openmp
torch.backends.openmp.is_available()


torch.backends.openmp.is_available()
Returns whether PyTorch is built with OpenMP support.
torch.backends.opt_einsum
torch.backends.opt_einsum.is_available()
Returns a bool indicating if opt_einsum is currently available.
Return type:
      bool
torch.backends.opt_einsum.get_opt_einsum()
Returns the opt_einsum package if opt_einsum is currently
   available, else None.
Return type:
      Any
torch.backends.opt_einsum.enabled
A :class:"bool" that controls whether opt_einsum is enabled ("True"
   by default). If so, torch.einsum will use opt_einsum (https
   ://optimized-einsum.readthedocs.io/en/stable/path_finding.html) if
   available to calculate an optimal path of contraction for faster
   performance.
If opt_einsum is not available, torch.einsum will fall back to the
   default contraction path of left to right.
torch.backends.opt_einsum.strategy
A :class:"str" that specifies which strategies to try when


"torch.backends.opt_einsum.enabled" is "True". By default,
   torch.einsum will try the "auto" strategy, but the "greedy" and
   "optimal" strategies are also supported. Note that the "optimal"
   strategy is factorial on the number of inputs as it tries all
   possible paths. See more details in opt_einsum's docs (https
   ://optimized-einsum.readthedocs.io/en/stable/path_finding.html).
torch.backends.xeon


torch.utils.dlpack
torch.utils.dlpack.from_dlpack(ext_tensor) -> Tensor
Converts a tensor from an external library into a "torch.Tensor".
The returned PyTorch tensor will share the memory with the input
   tensor (which may have come from another library). Note that in-
   place operations will therefore also affect the data of the input
   tensor. This may lead to unexpected issues (e.g., other libraries
   may have read-only flags or immutable data structures), so the user
   should only do this if they know for sure that this is fine.
Parameters:
      ext_tensor (object with "dlpack" attribute, or a DLPack
      capsule) --
  The tensor or DLPack capsule to convert.

  If "ext_tensor" is a tensor (or ndarray) object, it must support
  the "__dlpack__" protocol (i.e., have a "ext_tensor.__dlpack__"
  method). Otherwise "ext_tensor" may be a DLPack capsule, which
  is an opaque "PyCapsule" instance, typically produced by a


"to_dlpack" function or method.
Return type:
      Tensor
Examples:
  >>> import torch.utils.dlpack
  >>> t = torch.arange(4)

  # Convert a tensor directly (supported in PyTorch >= 1.10)
  >>> t2 = torch.from_dlpack(t)
  >>> t2[:2] = -1  # show that memory is shared
  >>> t2
  tensor([-1, -1,  2,  3])
  >>> t
  tensor([-1, -1,  2,  3])

  # The old-style DLPack usage, with an intermediate capsule object
  >>> capsule = torch.utils.dlpack.to_dlpack(t)
  >>> capsule
  <capsule object "dltensor" at ...>
  >>> t3 = torch.from_dlpack(capsule)
  >>> t3
  tensor([-1, -1,  2,  3])
  >>> t3[0] = -9  # now we're sharing memory between 3 tensors
  >>> t3
  tensor([-9, -1,  2,  3])
  >>> t2
  tensor([-9, -1,  2,  3])
  >>> t
  tensor([-9, -1,  2,  3])

torch.utils.dlpack.to_dlpack(tensor) -> PyCapsule
Returns an opaque object (a "DLPack capsule") representing the
   tensor.
Note:


tensor.
Note:
 "to_dlpack" is a legacy DLPack interface. The capsule it returns
 cannot be used for anything in Python other than use it as input
 to "from_dlpack". The more idiomatic use of DLPack is to call
 "from_dlpack" directly on the tensor object - this works when
 that object has a "__dlpack__" method, which PyTorch and most
 other libraries indeed have now.

Warning:
 Only call "from_dlpack" once per capsule produced with
 "to_dlpack". Behavior when a capsule is consumed multiple times
 is undefined.

Parameters:
      tensor -- a tensor to be exported
The DLPack capsule shares the tensor's memory.


PyTorch Governance | Build + CI
How to Add a New Maintainer
For the person to be a maintainer, a person needs to:

Land at least six commits to the related part of the PyTorch
  repository

At least one of these commits must be submitted in the last six
  months

To add a qualified person to the maintainers' list, please create a PR
that adds a person to the persons of interests page and merge_rules
files. Current maintainers will cast their votes of support. Decision
criteria for approving the PR: * Not earlier than two business days
passed before merging (ensure the majority of the contributors have
seen it) * PR has the correct label (module: ci) * There are no
objections from the current maintainers * There are at least three net
thumbs up from current maintainers (or all maintainers vote thumbs
up when the module has less than 3 maintainers).


Probability distributions - torch.distributions
The "distributions" package contains parameterizable probability
distributions and sampling functions. This allows the construction of
stochastic computation graphs and stochastic gradient estimators for
optimization. This package generally follows the design of the
TensorFlow Distributions package.
It is not possible to directly backpropagate through random samples.
However, there are two main methods for creating surrogate functions
that can be backpropagated through. These are the score function
estimator/likelihood ratio estimator/REINFORCE and the pathwise
derivative estimator. REINFORCE is commonly seen as the basis for
policy gradient methods in reinforcement learning, and the pathwise
derivative estimator is commonly seen in the reparameterization trick
in variational autoencoders. Whilst the score function only requires
the value of samples f(x), the pathwise derivative requires the


derivative f'(x). The next sections discuss these two in a
reinforcement learning example. For more details see Gradient
Estimation Using Stochastic Computation Graphs .
Score function
When the probability density function is differentiable with respect
to its parameters, we only need "sample()" and "log_prob()" to
implement REINFORCE:
\Delta\theta  = \alpha r \frac{\partial\log
   p(a|\pi^\theta(s))}{\partial\theta}
where \theta are the parameters, \alpha is the learning rate, r is the
reward and p(a|\pi^\theta(s)) is the probability of taking action a in
state s given policy \pi^\theta.
In practice we would sample an action from the output of a network,
apply this action in an environment, and then use "log_prob" to
construct an equivalent loss function. Note that we use a negative
because optimizers use gradient descent, whilst the rule above assumes
gradient ascent. With a categorical policy, the code for implementing
REINFORCE would be as follows:
probs = policy_network(state)


probs = policy_network(state)
   # Note that this is equivalent to what used to be called multinomial
   m = Categorical(probs)
   action = m.sample()
   next_state, reward = env.step(action)
   loss = -m.log_prob(action) * reward
   loss.backward()
Pathwise derivative
The other way to implement these stochastic/policy gradients would be
to use the reparameterization trick from the "rsample()" method, where
the parameterized random variable can be constructed via a
parameterized deterministic function of a parameter-free random
variable. The reparameterized sample therefore becomes differentiable.
The code for implementing the pathwise derivative would be as follows:
params = policy_network(state)
   m = Normal(*params)
   # Any distribution with .has_rsample == True could work based on the application
   action = m.rsample()
   next_state, reward = env.step(action)  # Assuming that reward is differentiable
   loss = -reward
   loss.backward()
Distribution


loss.backward()
Distribution
class torch.distributions.distribution.Distribution(batch_shape=torch.Size([]), event_shape=torch.Size([]), validate_args=None)
Bases: "object"
Distribution is the abstract base class for probability
   distributions.
property arg_constraints: Dict[str, Constraint]
  Returns a dictionary from argument names to "Constraint" objects
  that should be satisfied by each argument of this distribution.
  Args that are not tensors need not appear in this dict.

property batch_shape: Size
  Returns the shape over which parameters are batched.

cdf(value)
  Returns the cumulative density/mass function evaluated at
  *value*.

  Parameters:
     **value** (*Tensor*) --

  Return type:
     *Tensor*

entropy()
  Returns entropy of distribution, batched over batch_shape.

  Returns:
     Tensor of shape batch_shape.

  Return type:
     *Tensor*

enumerate_support(expand=True)


enumerate_support(expand=True)
  Returns tensor containing all values supported by a discrete
  distribution. The result will enumerate over dimension 0, so the
  shape of the result will be *(cardinality,) + batch_shape +
  event_shape* (where *event_shape = ()* for univariate
  distributions).

  Note that this enumerates over all batched tensors in lock-step
  *[[0, 0], [1, 1], ...]*. With *expand=False*, enumeration
  happens along dim 0, but with the remaining batch dimensions
  being singleton dimensions, *[[0], [1], ..*.

  To iterate over the full Cartesian product use
  *itertools.product(m.enumerate_support())*.

  Parameters:
     **expand** (*bool*) -- whether to expand the support over the
     batch dims to match the distribution's *batch_shape*.

  Returns:
     Tensor iterating over dimension 0.

  Return type:
     *Tensor*

property event_shape: Size
  Returns the shape of a single sample (without batching).


expand(batch_shape, _instance=None)
  Returns a new distribution instance (or populates an existing
  instance provided by a derived class) with batch dimensions
  expanded to *batch_shape*. This method calls "expand" on the
  distribution's parameters. As such, this does not allocate new
  memory for the expanded distribution instance. Additionally,
  this does not repeat any args checking or parameter broadcasting
  in *__init__.py*, when an instance is first created.

  Parameters:
     * **batch_shape** (*torch.Size*) -- the desired expanded
       size.

     * **_instance** -- new instance provided by subclasses that
       need to override *.expand*.

  Returns:
     New distribution instance with batch dimensions expanded to
     *batch_size*.

icdf(value)
  Returns the inverse cumulative density/mass function evaluated
  at *value*.

  Parameters:
     **value** (*Tensor*) --

  Return type:
     *Tensor*


Return type:
         Tensor
log_prob(value)
  Returns the log of the probability density/mass function
  evaluated at *value*.

  Parameters:
     **value** (*Tensor*) --

  Return type:
     *Tensor*

property mean: Tensor
  Returns the mean of the distribution.

property mode: Tensor
  Returns the mode of the distribution.

perplexity()
  Returns perplexity of distribution, batched over batch_shape.

  Returns:
     Tensor of shape batch_shape.

  Return type:
     *Tensor*

rsample(sample_shape=torch.Size([]))
  Generates a sample_shape shaped reparameterized sample or
  sample_shape shaped batch of reparameterized samples if the
  distribution parameters are batched.

  Return type:
     *Tensor*

sample(sample_shape=torch.Size([]))
  Generates a sample_shape shaped sample or sample_shape shaped
  batch of samples if the distribution parameters are batched.

  Return type:


Return type:
         Tensor
sample_n(n)
  Generates n samples or n batches of samples if the distribution
  parameters are batched.

  Return type:
     *Tensor*

static set_default_validate_args(value)
  Sets whether validation is enabled or disabled.

  The default behavior mimics Python's "assert" statement:
  validation is on by default, but is disabled if Python is run in
  optimized mode (via "python -O"). Validation may be expensive,
  so you may want to disable it once a model is working.

  Parameters:
     **value** (*bool*) -- Whether to enable validation.

property stddev: Tensor
  Returns the standard deviation of the distribution.

property support: Optional[Any]
  Returns a "Constraint" object representing this distribution's
  support.

property variance: Tensor
  Returns the variance of the distribution.

ExponentialFamily


ExponentialFamily
class torch.distributions.exp_family.ExponentialFamily(batch_shape=torch.Size([]), event_shape=torch.Size([]), validate_args=None)
Bases: "Distribution"
ExponentialFamily is the abstract base class for probability
   distributions belonging to an exponential family, whose probability
   mass/density function has the form is defined below
  p_{F}(x; \theta) = \exp(\langle t(x), \theta\rangle - F(\theta)
  + k(x))

where \theta denotes the natural parameters, t(x) denotes the
   sufficient statistic, F(\theta) is the log normalizer function for
   a given family and k(x) is the carrier measure.
Note:
 This class is an intermediary between the *Distribution* class
 and distributions which belong to an exponential family mainly to
 check the correctness of the *.entropy()* and analytic KL
 divergence methods. We use this class to compute the entropy and
 KL divergence using the AD framework and Bregman divergences


(courtesy of: Frank Nielsen and Richard Nock, Entropies and
     Cross-entropies of Exponential Families).
entropy()
  Method to compute the entropy using Bregman divergence of the
  log normalizer.

Bernoulli
class torch.distributions.bernoulli.Bernoulli(probs=None, logits=None, validate_args=None)
Bases: "ExponentialFamily"
Creates a Bernoulli distribution parameterized by "probs" or
   "logits" (but not both).
Samples are binary (0 or 1). They take the value 1 with
   probability p and 0 with probability 1 - p.
Example:
  >>> m = Bernoulli(torch.tensor([0.3]))
  >>> m.sample()  # 30% chance 1; 70% chance 0
  tensor([ 0.])

Parameters:
      * probs (Number, Tensor) -- the probability of
        sampling 1
  * **logits** (*Number**, **Tensor*) -- the log-odds of sampling
    *1*

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}
entropy()


entropy()
enumerate_support(expand=True)
expand(batch_shape, _instance=None)
has_enumerate_support = True
log_prob(value)
property logits
property mean
property mode
property param_shape
property probs
sample(sample_shape=torch.Size([]))
support = Boolean()
property variance
Beta
class torch.distributions.beta.Beta(concentration1, concentration0, validate_args=None)
Bases: "ExponentialFamily"
Beta distribution parameterized by "concentration1" and
   "concentration0".
Example:
  >>> m = Beta(torch.tensor([0.5]), torch.tensor([0.5]))
  >>> m.sample()  # Beta distributed with concentration concentration1 and concentration0
  tensor([ 0.1046])

Parameters:
      * concentration1 (float or Tensor) -- 1st
        concentration parameter of the distribution (often referred to
        as alpha)
  * **concentration0** (*float** or **Tensor*) -- 2nd


concentration parameter of the distribution (often referred to
        as beta)
arg_constraints = {'concentration0': GreaterThan(lower_bound=0.0), 'concentration1': GreaterThan(lower_bound=0.0)}
property concentration0
property concentration1
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
log_prob(value)
property mean
property mode
rsample(sample_shape=())
support = Interval(lower_bound=0.0, upper_bound=1.0)
property variance
Binomial
class torch.distributions.binomial.Binomial(total_count=1, probs=None, logits=None, validate_args=None)
Bases: "Distribution"
Creates a Binomial distribution parameterized by "total_count" and
   either "probs" or "logits" (but not both). "total_count" must be
   broadcastable with "probs"/"logits".
Example:
  >>> m = Binomial(100, torch.tensor([0 , .2, .8, 1]))
  >>> x = m.sample()
  tensor([   0.,   22.,   71.,  100.])


tensor([   0.,   22.,   71.,  100.])
  >>> m = Binomial(torch.tensor([[5.], [10.]]), torch.tensor([0.5, 0.8]))
  >>> x = m.sample()
  tensor([[ 4.,  5.],
          [ 7.,  6.]])

Parameters:
      * total_count (int or Tensor) -- number of Bernoulli
        trials
  * **probs** (*Tensor*) -- Event probabilities

  * **logits** (*Tensor*) -- Event log-odds

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0), 'total_count': IntegerGreaterThan(lower_bound=0)}
entropy()
enumerate_support(expand=True)
expand(batch_shape, _instance=None)
has_enumerate_support = True
log_prob(value)
property logits
property mean
property mode
property param_shape
property probs
sample(sample_shape=torch.Size([]))
property support
property variance
Categorical
class torch.distributions.categorical.Categorical(probs=None, logits=None, validate_args=None)
Bases: "Distribution"


Bases: "Distribution"
Creates a categorical distribution parameterized by either "probs"
   or "logits" (but not both).
Note:
 It is equivalent to the distribution that "torch.multinomial()"
 samples from.

Samples are integers from {0, \ldots, K-1} where K is
   "probs.size(-1)".
If probs is 1-dimensional with length-K, each element is the
   relative probability of sampling the class at that index.
If probs is N-dimensional, the first N-1 dimensions are treated
   as a batch of relative probability vectors.
Note:
 The *probs* argument must be non-negative, finite and have a non-
 zero sum, and it will be normalized to sum to 1 along the last
 dimension. "probs" will return this normalized value. The
 *logits* argument will be interpreted as unnormalized log
 probabilities and can therefore be any real number. It will
 likewise be normalized so that the resulting probabilities sum to


1 along the last dimension. "logits" will return this normalized
     value.
See also: "torch.multinomial()"
Example:
  >>> m = Categorical(torch.tensor([ 0.25, 0.25, 0.25, 0.25 ]))
  >>> m.sample()  # equal probability of 0, 1, 2, 3
  tensor(3)

Parameters:
      * probs (Tensor) -- event probabilities
  * **logits** (*Tensor*) -- event log probabilities
    (unnormalized)

arg_constraints = {'logits': IndependentConstraint(Real(), 1), 'probs': Simplex()}
entropy()
enumerate_support(expand=True)
expand(batch_shape, _instance=None)
has_enumerate_support = True
log_prob(value)
property logits
property mean
property mode
property param_shape
property probs
sample(sample_shape=torch.Size([]))
property support
property variance
Cauchy
class torch.distributions.cauchy.Cauchy(loc, scale, validate_args=None)
Bases: "Distribution"
Samples from a Cauchy (Lorentz) distribution. The distribution of


the ratio of independent normally distributed random variables with
   means 0 follows a Cauchy distribution.
Example:
  >>> m = Cauchy(torch.tensor([0.0]), torch.tensor([1.0]))
  >>> m.sample()  # sample from a Cauchy distribution with loc=0 and scale=1
  tensor([ 2.3214])

Parameters:
      * loc (float or Tensor) -- mode or median of the
        distribution.
  * **scale** (*float** or **Tensor*) -- half width at half
    maximum.

arg_constraints = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}
cdf(value)
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
icdf(value)
log_prob(value)
property mean
property mode
rsample(sample_shape=torch.Size([]))
support = Real()
property variance
Chi2
class torch.distributions.chi2.Chi2(df, validate_args=None)
Bases: "Gamma"
Creates a Chi-squared distribution parameterized by shape parameter


"df". This is exactly equivalent to "Gamma(alpha=0.5*df, beta=0.5)"
Example:
  >>> m = Chi2(torch.tensor([1.0]))
  >>> m.sample()  # Chi2 distributed with shape df=1
  tensor([ 0.1046])

Parameters:
      df (float or Tensor) -- shape parameter of the
      distribution
arg_constraints = {'df': GreaterThan(lower_bound=0.0)}
property df
expand(batch_shape, _instance=None)
ContinuousBernoulli
class torch.distributions.continuous_bernoulli.ContinuousBernoulli(probs=None, logits=None, lims=(0.499, 0.501), validate_args=None)
Bases: "ExponentialFamily"
Creates a continuous Bernoulli distribution parameterized by
   "probs" or "logits" (but not both).
The distribution is supported in [0, 1] and parameterized by
   'probs' (in (0,1)) or 'logits' (real-valued). Note that, unlike the
   Bernoulli, 'probs' does not correspond to a probability and
   'logits' does not correspond to log-odds, but the same names are


used due to the similarity with the Bernoulli. See [1] for more
   details.
Example:
  >>> m = ContinuousBernoulli(torch.tensor([0.3]))
  >>> m.sample()
  tensor([ 0.2538])

Parameters:
      * probs (Number, Tensor) -- (0,1) valued parameters
  * **logits** (*Number**, **Tensor*) -- real valued parameters
    whose sigmoid matches 'probs'

[1] The continuous Bernoulli: fixing a pervasive error in
   variational autoencoders, Loaiza-Ganem G and Cunningham JP, NeurIPS
   2019. https://arxiv.org/abs/1907.06845
arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}
cdf(value)
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
icdf(value)
log_prob(value)
property logits
property mean
property param_shape
property probs
rsample(sample_shape=torch.Size([]))
sample(sample_shape=torch.Size([]))
property stddev
support = Interval(lower_bound=0.0, upper_bound=1.0)


property variance
Dirichlet
class torch.distributions.dirichlet.Dirichlet(concentration, validate_args=None)
Bases: "ExponentialFamily"
Creates a Dirichlet distribution parameterized by concentration
   "concentration".
Example:
  >>> m = Dirichlet(torch.tensor([0.5, 0.5]))
  >>> m.sample()  # Dirichlet distributed with concentration [0.5, 0.5]
  tensor([ 0.1046,  0.8954])

Parameters:
      concentration (Tensor) -- concentration parameter of the
      distribution (often referred to as alpha)
arg_constraints = {'concentration': IndependentConstraint(GreaterThan(lower_bound=0.0), 1)}
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
log_prob(value)
property mean
property mode
rsample(sample_shape=())
support = Simplex()
property variance
Exponential
class torch.distributions.exponential.Exponential(rate, validate_args=None)
Bases: "ExponentialFamily"


Bases: "ExponentialFamily"
Creates a Exponential distribution parameterized by "rate".
Example:
  >>> m = Exponential(torch.tensor([1.0]))
  >>> m.sample()  # Exponential distributed with rate=1
  tensor([ 0.1046])

Parameters:
      rate (float or Tensor) -- rate = 1 / scale of the
      distribution
arg_constraints = {'rate': GreaterThan(lower_bound=0.0)}
cdf(value)
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
icdf(value)
log_prob(value)
property mean
property mode
rsample(sample_shape=torch.Size([]))
property stddev
support = GreaterThanEq(lower_bound=0.0)
property variance
FisherSnedecor
class torch.distributions.fishersnedecor.FisherSnedecor(df1, df2, validate_args=None)
Bases: "Distribution"
Creates a Fisher-Snedecor distribution parameterized by "df1" and
   "df2".
Example:
  >>> m = FisherSnedecor(torch.tensor([1.0]), torch.tensor([2.0]))


m.sample()  # Fisher-Snedecor-distributed with df1=1 and df2=2
      tensor([ 0.2453])

Parameters:
      * df1 (float or Tensor) -- degrees of freedom
        parameter 1
  * **df2** (*float** or **Tensor*) -- degrees of freedom
    parameter 2

arg_constraints = {'df1': GreaterThan(lower_bound=0.0), 'df2': GreaterThan(lower_bound=0.0)}
expand(batch_shape, _instance=None)
has_rsample = True
log_prob(value)
property mean
property mode
rsample(sample_shape=torch.Size([]))
support = GreaterThan(lower_bound=0.0)
property variance
Gamma
class torch.distributions.gamma.Gamma(concentration, rate, validate_args=None)
Bases: "ExponentialFamily"
Creates a Gamma distribution parameterized by shape "concentration"
   and "rate".
Example:
  >>> m = Gamma(torch.tensor([1.0]), torch.tensor([1.0]))
  >>> m.sample()  # Gamma distributed with concentration=1 and rate=1
  tensor([ 0.1046])

Parameters:


tensor([ 0.1046])
Parameters:
      * concentration (float or Tensor) -- shape parameter
        of the distribution (often referred to as alpha)
  * **rate** (*float** or **Tensor*) -- rate = 1 / scale of the
    distribution (often referred to as beta)

arg_constraints = {'concentration': GreaterThan(lower_bound=0.0), 'rate': GreaterThan(lower_bound=0.0)}
cdf(value)
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
log_prob(value)
property mean
property mode
rsample(sample_shape=torch.Size([]))
support = GreaterThanEq(lower_bound=0.0)
property variance
Geometric
class torch.distributions.geometric.Geometric(probs=None, logits=None, validate_args=None)
Bases: "Distribution"
Creates a Geometric distribution parameterized by "probs", where
   "probs" is the probability of success of Bernoulli trials. It
   represents the probability that in k + 1 Bernoulli trials, the


first k trials failed, before seeing a success.
Samples are non-negative integers [0, \inf).
Example:
  >>> m = Geometric(torch.tensor([0.3]))
  >>> m.sample()  # underlying Bernoulli has 30% chance 1; 70% chance 0
  tensor([ 2.])

Parameters:
      * probs (Number, Tensor) -- the probability of
        sampling 1. Must be in range (0, 1]
  * **logits** (*Number**, **Tensor*) -- the log-odds of sampling
    *1*.

arg_constraints = {'logits': Real(), 'probs': Interval(lower_bound=0.0, upper_bound=1.0)}
entropy()
expand(batch_shape, _instance=None)
log_prob(value)
property logits
property mean
property mode
property probs
sample(sample_shape=torch.Size([]))
support = IntegerGreaterThan(lower_bound=0)
property variance
Gumbel
class torch.distributions.gumbel.Gumbel(loc, scale, validate_args=None)
Bases: "TransformedDistribution"
Samples from a Gumbel Distribution.
Examples:


Examples:
  >>> m = Gumbel(torch.tensor([1.0]), torch.tensor([2.0]))
  >>> m.sample()  # sample from Gumbel distribution with loc=1, scale=2
  tensor([ 1.0124])

Parameters:
      * loc (float or Tensor) -- Location parameter of the
        distribution
  * **scale** (*float** or **Tensor*) -- Scale parameter of the
    distribution

arg_constraints: Dict[str, constraints.Constraint] = {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0)}
entropy()
expand(batch_shape, _instance=None)
log_prob(value)
property mean
property mode
property stddev
support = Real()
property variance
HalfCauchy
class torch.distributions.half_cauchy.HalfCauchy(scale, validate_args=None)
Bases: "TransformedDistribution"
Creates a half-Cauchy distribution parameterized by scale where:
  X ~ Cauchy(0, scale)
  Y = |X| ~ HalfCauchy(scale)

Example:
  >>> m = HalfCauchy(torch.tensor([1.0]))


m = HalfCauchy(torch.tensor([1.0]))
      >>> m.sample()  # half-cauchy distributed with scale=1
      tensor([ 2.3214])

Parameters:
      scale (float or Tensor) -- scale of the full Cauchy
      distribution
arg_constraints: Dict[str, constraints.Constraint] = {'scale': GreaterThan(lower_bound=0.0)}
cdf(value)
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
icdf(prob)
log_prob(value)
property mean
property mode
property scale
support = GreaterThanEq(lower_bound=0.0)
property variance
HalfNormal
class torch.distributions.half_normal.HalfNormal(scale, validate_args=None)
Bases: "TransformedDistribution"
Creates a half-normal distribution parameterized by scale where:
  X ~ Normal(0, scale)
  Y = |X| ~ HalfNormal(scale)

Example:
  >>> m = HalfNormal(torch.tensor([1.0]))
  >>> m.sample()  # half-normal distributed with scale=1
  tensor([ 0.1046])

Parameters:


tensor([ 0.1046])
Parameters:
      scale (float or Tensor) -- scale of the full Normal
      distribution
arg_constraints: Dict[str, constraints.Constraint] = {'scale': GreaterThan(lower_bound=0.0)}
cdf(value)
entropy()
expand(batch_shape, _instance=None)
has_rsample = True
icdf(prob)
log_prob(value)
property mean
property mode
property scale
support = GreaterThanEq(lower_bound=0.0)
property variance
Independent
class torch.distributions.independent.Independent(base_distribution, reinterpreted_batch_ndims, validate_args=None)
Bases: "Distribution"
Reinterprets some of the batch dims of a distribution as event
   dims.
This is mainly useful for changing the shape of the result of
   "log_prob()". For example to create a diagonal Normal distribution
   with the same shape as a Multivariate Normal distribution (so they
   are interchangeable), you can: