How the autograd is implmented in pytorch

moha657 · April 13, 2025, 10:25am

Hi,

I am currently trying to understand the PyTorch codebase. For now, the implementations of the Linear layer, for example, are described by these two files in GitHub repos, but I can’t understand how the operations are stored for the computational graph.

github.com/pytorch/pytorch

torch/nn/modules/linear.py

v2.6.0


      
          
              """
          
              def __init__(self, *args: Any, **kwargs: Any) -> None:
                  super().__init__()
          
              def forward(self, input: Tensor) -> Tensor:
                  return input
          
          
          class Linear(Module):
              r"""Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`.
          
              This module supports :ref:`TensorFloat32<tf32_on_ampere>`.
          
              On certain ROCm devices, when using float16 inputs this module will use :ref:`different precision<fp16_on_mi200>` for backward.
          
              Args:
                  in_features: size of each input sample
                  out_features: size of each output sample
                  bias: If set to ``False``, the layer will not learn an additive bias.

github.com/pytorch/pytorch

torch/csrc/api/src/nn/modules/linear.cpp

main

#include <torch/nn/functional/linear.h>
#include <torch/nn/init.h>
#include <torch/nn/modules/linear.h>

#include <torch/types.h>
#include <torch/utils.h>

#include <cmath>
#include <cstdint>

namespace F = torch::nn::functional;

namespace torch::nn {

void IdentityImpl::reset() {}

void IdentityImpl::pretty_print(std::ostream& stream) const {
  stream << "torch::nn::Identity()";
}

This file has been truncated. show original

soulitzer · April 13, 2025, 10:22pm

You might be interested in
https://blog.ezyang.com/2019/05/pytorch-internals/
http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/

soulitzer · April 13, 2025, 10:23pm

Short answer is that there will be some additional logic that interposes to track the operations everytime you run an operation.

moha657 · April 13, 2025, 10:38pm

Thanks for the response the articles are very interesting in fact what i am interested in is knowing at a deep level the in and out of pytorch as i wan to go from a intermediate level into upper intermediate in both pytorch and deep learning