# How to read the autograd codebase

**URL:** <https://dev-discuss.pytorch.org/t/how-to-read-the-autograd-codebase/383>\
**Category:** frontend API\
**Created:** [October 26, 2021, 3:36pm UTC](https://dev-discuss.pytorch.org/t/how-to-read-the-autograd-codebase/383 "2021-10-26T15:36:05Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![albanD](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/alband/32/9_2.png) [@albanD](https://dev-discuss.pytorch.org/u/albanD)\
**Post date:** [October 26, 2021, 3:36pm UTC](https://dev-discuss.pytorch.org/t/how-to-read-the-autograd-codebase/383/1 "2021-10-26T15:36:06Z")

</div>

# How to read the autograd code in PyTorch

This document will try to give you a good idea of how to browse the autograd-related source in PyTorch The goal is to get you familiar with what the key pieces are, where they are located, and the order in which you should read them.  
Warning - this is by no means trying to give a good example of how to do things but a current state. In retrospect, we could do quite a few things better.

Note that all the file locations and links are done at [ddf2681](https://github.com/pytorch/pytorch/tree/ddf26816d3ca54ce7f3513f618fac93ce67d06e9). Things might have moved in master since this has been written. Let me know if this gets outdated and I can update the links. At the time of writing, the forward mode AD code is not complete, so you can ignore all reference to it for now. The profiler code will also be ignored in this discussion as it is fairly independent from the autograd and maintained by different people.

This post assumes that you are familiar with Automatic Differentiation and what it tries to do.  
All the discussion here will be about backward mode AD (or back propagation). We do backward mode AD by creating a computational graph while the function is being evaluated (forward pass) and then executing this graph when the user asks for gradients (backward pass).

At a very high level, autograd needs to provide the following functionalities:

- Build computational graph along with the forward
- Evaluate that graph to do a backward
- Provide high level APIs
- Provide testing tools

The short version of the full description below is

- Low level data structures that are the components of the computational graph (Node, SavedVariable, Tensor)
- Wrappers for forward functions to build the graph
  - Manually in C++
  - Via codegen and derivatives.yaml
  - Via custom Function (in Python or C++)
  - Special care for view/inplace

- An engine to execute the computational graph
- Hook system to control execution of the graph
- Python bindings for all c++ objects
- A Python and a C++ API
- A high level API to compute complex quantities in Python
- Gradient checking utilities for testing
- Anomaly detection utility for debugging and testing
- Dispatcher and how autograd integrates into PyTorch

## Basic Components of the Computational Graph

**The Node class**  
The graph that we build is a DAG based on the Node class. You can read its description [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/function.h#L50-L97).  
You can also read the rest of the Node class implementation in `csrc/autograd/function.`{`h,cpp}` with special care to:

- `operator()()` that is used to evaluate this Node
- `next_edges()` that is used to find the child Nodes in the graph (correspond to the `.next_functions` attribute of the python binding)
- `sequence_nr()` that is used to have a, thread local, universal ordering of the Node. This ordering is used to ensure determinism in the execution order of the engine (when only a single device is used).
- `pyobj()` that is used to store the corresponding python wrapper
- `metadata()` that is used to store user metadata associated with this Node, see [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/anomaly_mode.h#L52-L61) for the c++ metadata and [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_anomaly_mode.h#L10-L32) for the python metadata.
- `release_variables()` that is used to release all the resources saved by that Node

**Subclasses of Node**  
We can’t directly create instances of Node because it is abstract. All nodes in our DAG are actually instances of Node’s _subclasses_. These usually implement at least:

- A custom constructor
- An `apply()` function that is called when the Node is evaluated
- `release_variables()` if this Node store any states at constructions

Since the backward often needs some values from the forward to perform its job, subclasses of Node also typically store extra fields of all the inputs/outputs from the forward computation that it needs to use to compute the backward pass, we have a a special class to save these (without creating ref cycles that would leak memory): SavedVariables.

**Saved Variables**  
SavedVariable is implemented in `csrc/autograd/saved_variable.{h,cpp}`  
There are only two functions there:

- The constructor that saves the plain Tensor data, stores the fields of autograd metadata on the side (this breaks the cycle mentioned above) and save the current “version” of the Tensor.
- The `.unpack()` function that can be used to get the original Tensor + autograd metadata back. It also checks that the “version” of the Tensor didn’t change to make sure that no in-place operation was done on it since it was saved.

_What would happen if we didn’t have SavedVariables?_  
Note that the cycle mentioned above would appear because the Node saves the Tensor while the Tensor has a `grad_fn` field which link back to the Node. This class avoids this cycle by making sure that the saved Tensor is not the same as the one that has the `grad_fn` field.

EDIT: More work has happening in this class since then. In particular, there are now hooks that can be used to specify what should happen while packing/unpacking. See [this doc](https://pytorch.org/docs/stable/notes/autograd.html#hooks-for-saved-tensors) for more details.

**AutogradMeta**  
In addition to its grad\_fn, the output tensor’s “entry point” to the backward graph we described above, every tensor also holds other metadata that are important to autograd. This metadata is stored in the AutogradMeta struct which is uniquely owned by each tensor. In addition to a reference to grad\_fn\_ (owning), AutogradMeta also stores: reference to grad\_accumulator\_ (weak), a reference to grad\_ Tensor (owning), as well as some other fields.

As of today in PyTorch, all Tensors are autograd aware and can store such metadata as seen [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/c10/core/TensorImpl.h#L1678-L1702). This used to be different and we had `Variable`s that were the autograd-aware wrapper for Tensors. As you can see [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/variable.h#L25-L32), the two are exactly the same type now and you can read `Variable` in the autograd codebase as `Tensor`.  
Also, we have a special kind of Node associated with every Tensor (in a lazy manner as seen [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/variable.cpp#L217-L229)) that is called AccumulateGrad. These are the leaf Nodes in the DAG and will, when executed, accumulate gradients in the `.grad` field of the corresponding Tensor.

That is all!  
Now it’s all about building this DAG properly and executing it.

## Building the Graph

This section presents how we build the graph and all the subtleties to make sure it computes the right thing in all case (or raises an error if we cannot handle a particular case).  
The next section will present how to execute this graph in details.

### Manually

The most basic way to build the graph can be found [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/functions/utils.cpp#L12-L38).  
This function takes the inputs of the computation, the output of the computations as well as a lambda that creates the Node and builds the graph accordingly. In particular, it:

- If no input require gradients, make all input requires\_grad=false and return
- Collect the next edges from the inputs to the function
- Build the Node based on these next edges
- For every output Tensor, make it point properly to the newly created Node

We have a (very) small number of Nodes that are created this way and are implemented in the `csrc/autograd/functions/` folder. All of them are here because they require very manual tuning of what happens that does not fit with the more generic tools we provide to create the graph.  
We will come back to most of these but a simple example is the `Error` Node (that fails in the forward) that will just raise an error if anyone tries to use it. You can see the definition in [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/functions/basic_ops.h#L13-L24) and the apply function’s implementation in [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/functions/basic_ops.cpp#L14-L16).  
Another example is the `DelayedError` one [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/functions/basic_ops.h#L38-L49) that will raise an error if the backward of this Node tries to be called.

### Codegen

The problem is that writing a new class for every single operator in PyTorch would be very tedious: introducing the codegen.

Indeed, we use `tools/autograd/derivatives.yaml` as a configuration file that specify the backward formulas for (almost) all our operators.

The codegen should:

- Replace the manual code above to define the Node`
- Add code to automatically create the graph along with the forward
- Add code to make factory functions (like ` torch.tensor()` or `torch.rand()`) autograd-aware

For this part, I would recommend to first readers to skip reading the python code altogether. Reading examples from the generated files will be enough to understand how things work. For a good example, you can look for the functions associated with `abs` and `angle`.  
A reading guide for the python version would be another post altogether and left for future work.

The main entry point for the codegen is `tools/autograd/gen_autograd.py` that will generate all the necessary files in `torch/csrc/autograd/generated/`. In particular, it will use:

- `tools/autograd/gen_autograd_functions.py` to
  - generate all the Nodes that are required in `torch/csrc/autograd/generated/Functions.{h,cpp}`. It will make sure for each of them that:
    - The `.apply()` method performs the computation specified in `derivatives.yaml`
    - All inputs that will be used in the backward computation are saved and the `.release_variables()` method will release everything that was saved
    - The `.name()` method returns the name of this Node

  - generate python bindings for these Nodes in `torch/csrc/autograd/generated/python_functions.{h,cpp}` to allow to be accessed from python easily (and have checks like `isinstance` work properly).
  - You should look for `AbsBackward` and `AngleBackward` for the two examples.
  - Note as well that helper functions like `angle_backward()` are implemented in `torch/csrc/autograd/FunctionsManual.{h,cpp}`.

- `tools/autograd/gen_variable_type.py` to generate the code that wraps ALL the Tensor methods in PyTorch in `torch/csrc/autograd/generated/VariableTypeEverything.cpp`.
  - These wrappers are responsible for:
    - Checking if the output will require gradients
    - Creating an instance of the Node for that op if needed
    - Saving the relevant input/outputs
    - Building the graph accordingly
    - Setting up the view relationship
    - Perform in-place checks and version counter bumps

  - You can look for `abs()` and `angle()` in this file
  - Note that some functions have special wrappers that can be found in `torch/csrc/autograd/VariableTypeManual.{h,cpp}`. These are used to implement functions that are not supported by the codegen or do not compute derivatives (like detach).

- `tools/autograd/gen_trace_type.py` that is very similar to variable type above but it generates a file called `TraceTypeEverything.cpp` and is only responsible to perform tracing if it is enabled.
- `tools/autograd/gen_variable_factories.py` that generates ``torch/csrc/autograd/generated/variable\_factories.h`. This contains a wrapper for each factory function that makes sure that the `requires\_grad=`argument is properly handled. You can check`zeros()` in there for an example.
- `tools/autograd/load_derivatives.py` is a helper file to interpret the content of `derivatives.yaml`, nothing to read here.

### Custom Functions

While the codegen allows us to generate all this code very easily, it is not a convenient tool to use for third party users that want to define new Nodes: introducing custom Functions.

These allow the user to specify for the forward and backward of a given operation and will do all the necessary work to build the graph and setup the proper Node.  
A good introduction on how to use the python API for this is here: [Extending PyTorch — PyTorch 2.9 documentation](https://pytorch.org/docs/stable/notes/extending.html) and the c++ API mimics it as much as possible.

The C++ implementation is the simplest and is done with:

- You can find [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/custom_function.h#L79-L88) the base class that users should extend and define the forward and backward for.
- And [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/custom_function.h#L204-L248) is the `.apply()` function called during the forward that is responsible for calling the user’s forward as well as setting up the graph.
- The most important function here is the `._wrap_outputs()` one [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/custom_function.cpp#L20-L158) that is used both by the python and c++ API and is the one setting up the graph. You should take the time to read this function in detail to understand what it is doing.
- This Node is of a special type defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/custom_function.h#L151-L167) that is stores the necessary info as well as the AutogradContext, defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/custom_function.h#L90-L137), that allows to reproduce a similar API to the python one.
- This Node’s `.apply()` method [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/custom_function.h#L250-L313) is then responsible for calling the user defined backward as well as releasing the related resources

The Python implementation is more complex as most of the components are still implemented in c++ using CPython C API:

- When reading this code, you can ignore all the code related to instantiating or calling the Function class instances. This is the old API for custom Function that will be removed soon.
- The user facing python class can be found [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/autograd/function.py#L121). It defines the two functions that should be overwritten by the user.
- You can check the `_ContextMethodMixin` [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/autograd/function.py#L12) that defines all the functions available on the `ctx` for the python users.
- The `_HookMixin` [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/autograd/function.py#L73) allows the user to register hooks on these custom Functions (not sure if this one is actually documented).
- The `FunctionMeta` [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/autograd/function.py#L92) that handles the fallback to old style Functions
- The `_FunctionBase` that is defined in c++ [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_function.cpp#L1027-L1066) it implements all the low level API associated with the `ctx`
  - You can see accessors for all the attributes saved on the c++ Node [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_function.cpp#L886-L1016)
  - You can see its `.apply()` method [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_function.cpp#L621-L677) that is responsible for creating the ctx, calling the user specified forward and setting up the graph properly with the same `._wrap_outputs()` function as above (be careful we have two functions with that name that call each other!)
  - The Node it adds in the graph is a PyNode defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_function.h#L22-L55) that has a special apply function [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_function.cpp#L98-L189) that is responsible for calling the python class’s backward via the `.apply()` of the python class defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/autograd/function.py#L87-L89). The PyNode apply is also responsible for packing inputs into python objects and unpacking the results back into c++ objects.

### Special care for view and in-place

PyTorch’s autograd allows the user to use views and in-place operations but special care needs to be taken to make these works.  
The gist of what needs to be done here is:

- Ensure that all differentiable views are properly tracked.
- Ensure that in-place operation properly bump version counters.
- Ensure that in-place operations properly do graph rewrites if needed.
- Ensure that views get the right graph if their base or another view was modified in-place.

The view tracking is done via the `as_view()` function [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/VariableTypeUtils.h#L136-L245). This function is called by all the codegen’ed and manual functions if needed. You can check the codegen for ` view()` for example in `VariableTypeEverything.cpp` to see that call. You can also check the manual code for `detach()` in `VariableTypeManual.cpp` for an example of a non-differentiable view.  
You can then read the `make_variable_differentiable_view()` in `torch/csrc/autograd/variable.h` to see how this is done: namely using a special type of autograd metadata that also store the view informations.  
Note that custom Function do not do this! Indeed, the view info are tracked even in no\_grad mode and so we can simply read the view info on the output fo the custom Function to know if it is a view. This won’t work if the user reads the raw content of the Tensor in the forward and create a new Tensor that alias it but we are happy with this limitation.

The version counter bump is done in the codegen as well via the `.increment_version()` function call. You can find it for example when looking for `relu_()` in `VariableTypeEverything.cpp`.  
The custom functions also do this via the `.mark_dirty()` function on the context. It allows the user to tell the autograd what they modified in-place.

When an in-place operation happens, we need to rewrite the graph to make sure that all the views of that same memory see this change in their history. This is done by the `.rebase_history()` function that is called both by the codegen and the custom Functions if needed. The main entry point is [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/VariableTypeUtils.h#L95-L115) and calls into the c++ implementation [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/variable.cpp#L134-L164).  
As you can see there, this rewrite is only needed if the Tensor being modified in-place is a view. In that case, the following happen:

- Get the base of that view.
- Generate a new CopySlices Node that wraps the Node corresponding to the in-place op
  - This special Node expects a gradient of the same size as the base
  - Slice that gradient to be the size of the view being worked on
  - Call the original Node corresponding to the in-place op
  - Embed this small gradient in the full gradient buffer that is the same size as the base

- Set this CopySlices as the new grad\_fn for the base → meaning that this grad\_fn will now be used by all the views!
- Trigger an update of the grad\_fn for this view implemented [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/variable.cpp#L401-L471)
  - If this Tensor is a view and has been modified in-place since last time we generated its grad\_fn (checked via the “version”)
  - regenerate the graph between the base of this Tensor and this Tensor
  - Set the grad\_fn as this newly generated graph (that now points properly to the CopySlices generated above)

Ensuring that all other views properly point to the right graph is done automatically via the special grad\_fn call described above. Indeed, other views will have their “version” updated as well by the version counter bump and so it will trigger a recompute of their graph the first time someone tries to access it.

## Executing the Graph: The Engine

The execution engine is quite complex but fairly well contained. It is completely implemented in `torch/csrc/autograd/engine.{h,cpp}` and `torch/csrc/autograd/python_engine.{h,cpp}`.

We will get back to the python version of the engine at the end of this section as it only does very minimal changes.

The Engine is a stateful class for which a single instance exist in the whole PyTorch process. It is defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L984-L990) and is created along with the process.  
The high level ideas for the engine are:

- Given outputs and inputs of the computations, it should run all the graph to compute gradients for the inputs.
- It runs the backward on the same device as the forward
- It runs CPU backward on the thread that called the engine. All the GPU backward run on different worker threads.
- The engine is re-entrant as it may be called from within a backward Node. It must have measures to avoid stack overflows when there are deep recursive calls like this.
- The engine provides a callback system that allows the user to run a function once the execution of a given task (to be defined below) is finished.
- 

The engine has a single public API to execute the graph: `.execute()`.  
The only other user-facing API is the one to register callbacks. We will ignore this one for now.

You can read the code in the following order:

- Start reading the `.execute()` code [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.h#L268-L276).
  - An important data structure there is the GraphTask defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.h#L48-L178) that will be used to store everything that we need corresponding to this particular invocation of the execute function. In particular, it store things like
    - What remains to be done
    - If an error happened and if it is done
    - execute’s arguments such as keep\_graph and grad\_mode
    - Buffers for the gradients currently flowing back
    - Execution info about what needs to be executed for this particular call to execute. Remember that the same graph can be used (concurrently!) by multiple execute call. So the execute call must never modify the graph.
    - The calling thread TLS to be able to propagate it to worker threads
    - CUDA stream infos
    - re-entrant info
    - cpu work queue
    - callbacks

  - You can read all the functions called there, in particular `compute_dependencies` that is responsible for checking what needs to run in the graph
  - Ignore `execute_with_graph_task` for now
  - Note that at the end of this function, accessing the future’s value will throw if a worker thread sets an exception on the future. This is how the exception stack traces are propagated from the worker thread back to the main thread. Also this wait is always a no-op as, by this time, the execution is finished.

- Go into `execute_with_graph_task` now that can be found [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L910-L982)
  - It will set the final part of the graph task
  - Then push onto the right ready queue the first task that needs to be done
  - And it will go back to work
    - If it is a thread that was not involved in autograd before, start processing the CPU ready queue via `thread_main`. This function will exit once nothing remains to be done for this graph task
    - If this was already a worker thread (meaning we’re doing re-entrant autograd)
      - If you’re not too deep, update the depth and go into `thread_main`
      - If you’re too deep and risk to stack overflow, create a new thread via `add_thread_pool_task` that will do the work for you.

- You can now check `thread_main` [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L308-L410) that contains the main evaluation loop. You can follow all the calls into `evaluate_function` to see how we prepare the inputs, call the Node hooks, detect/set errors, validate the shapes, populate the input buffers, etc
- Once all the tasks are done and the future is marked as such, the post processing function [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L480-L513) is called as well to call the callbacks and do the cuda stream postprocessing.
- It is important to see as well that the `start_device_thread` function is called with an std::call\_once and is implemented [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L1063-L1093). It is used to initialize the worker threads (lazily the first time a user calls into the engine) for the engine singleton. You can see that each worker thread runs `thread_init` implemented [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L258-L294) and sets it’s device properly before calling `thread_main` with an empty graph\_task, meaning that it will run until he receives a special kill task.
- Similarly, the re-entrant threadpool has long running threads that will execute one graph\_task at a time as shown [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L412-L439).

The python engine has only very limited differences with the main engine:

- It has a special pair of thread\_on\_exception [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_engine.cpp#L69) and execute [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_engine.cpp#L84) to save the full python stack trace on the worker thread and restore it properly on the main thread.
- Special logic in the worker thread to pre-initialize the python state [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_engine.cpp#L59-L60) to make sure that GIL acquisition in the worker threads is fast.

## Hook system

As you’ve seen in the engine code, it is properly calling both pre and post hooks that are registered on the Node.

You can see these being attributes of all the Nodes in `torch/csrc/autograd/function.h` and you can find their definition in `torch/csrc/autograd/function_hook.{h,cpp}`.

You can also find specialization of these hooks to be able to register arbitrary `std::function` in `torch/csrc/autograd/cpp_hook.{h,cpp}`.  
And you can find specialization to be able to register python functions in `torch/csrc/autograd/python_hook.{h,cpp}`.

All of these should be self-contained and easy to read.

## Python bindings for all c++ objects

All the files in `torch/csrc/autograd/` that start with `python_` are actually implementing python-only logic.  
We’ve seen a couple above related to the engine and the hooks, the other ones mainly define the python bindings for all the objects we have in c++.

Variables (Tensors)  
A representative example here is the manual binding we do for Variable, which are actually Tensors.  
You can see the custom python object defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_variable.h#L10-L19).  
The corresponding python type is defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_variable.cpp#L755-L794). The functions provided in this struct are important as they should follow the very strict rules for the CPython type API. In particular, `THPVariable_pynew` and `THPVariable_dealloc` are custom creators and destructors. And `THPVariable_traverse` and `THPVariable_clear` are used to ensure ref cycles can be cleared by python’s GC. You can check these implementations if you’re curious.  
With all the special attributes and methods for it defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_variable.cpp#L704-L753). You can check above in this file for the implementation for all of these.  
You can then see [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_variable.cpp#L828) that this is the type that is exported as `_TensorBase` to the python side and that is used by torch as the parent class for all the Tensor types.

Nodes  
You can also find the python binding for the Node created for python custom Functions [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_function.h#L74-L116).  
The binding for all the other Nodes that only live in c++ is [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_cpp_function.h#L13-L16). This one is a bit special because we create a different python class for every Node we define in c++.  
The binding for the execution engine is [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_engine.cpp#L24-L26) and is very simple as it only exposes two functions (the one to run backward and the one to add callbacks).

(continuing in the reply below because of the word limit)

---

<div class="post-metadata">

**Author:** ![albanD](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/alband/32/9_2.png) [@albanD](https://dev-discuss.pytorch.org/u/albanD)\
**Post date:** [October 26, 2021, 3:36pm UTC](https://dev-discuss.pytorch.org/t/how-to-read-the-autograd-codebase/383/2 "2021-10-26T15:36:12Z")

</div>

## A Python and a C++ API

The end -user API in both cases is composed of two functions:

- backward() that runs the whole graph and accumulate gradient in all leaf’s .grad fields
- grad() that returns the gradients for the specified inputs

The end-user c++ API is implemented in `torch/csrc/autograd/autograd.{h,cpp}`.  
It is a very thin layer above the engine execute function. You can read these two files in details.

Note as well that the Tensor.backward() function in c++ is implemented in `VariableTypeManual.cpp` and is calling into this user-facing API directly.

The end-user python API is implemented in `torch/autograd/init.py` and should be fairly straightforward as well as it only does error checking and then calls into the c++ implementation.  
It calls into the `run_backward()` function [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/python_engine.cpp#L122-L253) that is unpacking the python objects, releasing the GIL and then calling into the engine.

## A high level API to compute complex quantities in Python

The functional API is a self contained API that builds on top of the basic API mentioned in the previous section.  
It is implemented in `torch/autograd/functional.py`.

You should read the functions there in the following order:

- vjp
- jacobian
- vhp
- hessian
- jvp
- hvp

## Gradient checking utilities for testing

In python, we provide some gradient checking utilities to make sure that the gradients computed by the autograd are correct.  
There are two functions for this `gradcheck()` and `gradgradcheck()` that can be found in `torch/autograd/gradcheck.py`.

You should read `gradcheck` first and follow each function call in that file.  
This function does:

- validation of the inputs
- make sure that if no output require gradients, the numerical gradients are all 0
- make sure that the analytical (computed via autograd) gradients match the numerical gradients. With special handling for the complex case
- make sure the backward pass is not ignoring the grad\_output that is given to it by passing a 0 grad\_output and making sure that all gradients are 0
- make sure that the gradient formula handle undefined Tensors properly (we use undefined Tensors to represent 0 gradients)

Then `gradgradcheck` is doing the same thing but for second order gradients.

## Anomaly detection utility for testing

Anomaly detection’s goal is to help user debug errors that happen in the backward pass.  
In particular, when bad things happen in the backward, since it runs (almost) completely in c++, it can be hard to know what is happening from the python side.  
To this end, the anomaly detection allows two things:

- detect nan in the gradient as soon as they appear so that you know directly which Node created the nan values first.
- when an error happen while executing a given Node in the backward, print also the stack trace of the forward when this Node was created.

The anomaly mode is enabled from python via a context manager defined in `torch/autograd/anomaly_mode.py`.  
It is enabled from c++ via a similar RAII guard defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/anomaly_mode.h#L24-L50).  
Both of these set a global flag that is defined [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/anomaly_mode.h#L12-L22) and [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/anomaly_mode.cpp#L10).

The nan detection is implemented in the engine [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L740-L751).

For the enhanced stack traces, it is a bit more complex:

- It hooks into every Node creation [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/function.h#L108-L116) and store the current stack trace (in a slightly different way if it is in c++ or python, see each implementation for details) as well as its parent (in case of re-entrant autograd, which Node triggered the autograd call that lead to the creation of this new Node).
- When an error happens during the execution of the engine, print the stack trace of the forward (and all its parents) [here](https://github.com/pytorch/pytorch/blob/ddf26816d3ca54ce7f3513f618fac93ce67d06e9/torch/csrc/autograd/engine.cpp#L517-L520).

That should give a good overview of where all the main pieces are. There are still some part that are missing like the role of the dispatcher or the new inference mode.  
Please leave any comment you have or question below!