2.6/_modules/torch/autograd/graph.html



<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta name="robots" content="noindex">
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>torch.autograd.graph &mdash; PyTorch 2.6 documentation</title>
  

    <link rel="canonical" href="https://pytorch.org/docs/stable/_modules/torch/autograd/graph.html"/>
  

  <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
  <!-- <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" /> -->
  <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/copybutton.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.10.0-beta/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.13.11/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/katex-math.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/sphinx-dropdown.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/panels-bootstrap.min.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/css/jit.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/css/custom.css" type="text/css" />
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" />


  <!-- Google Tag Manager -->
    <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start':
    new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0],
    j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src=
    'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f);
    })(window,document,'script','dataLayer','GTM-T8XT4PS');</script>
    <!-- End Google Tag Manager -->
  

  <script src="../../../_static/js/modernizr.min.js"></script>

  <!-- Preload the theme fonts -->

<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">

<!-- Preload the katex fonts -->

<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
  <link rel="stylesheet" href="https://use.fontawesome.com/releases/v5.15.2/css/all.css" integrity="sha384-vSIIfh2YWi9wW0r9iZe7RJPrKwp6bG+s9QZMoITbCckVJqGCCRhc+ccxNcdpHuYu" crossorigin="anonymous">
</head>

<div class="container-fluid header-holder tutorials-header" id="header-holder">
  <div class="container">
    <div class="header-container">
      <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>

      <div class="main-menu">
        <ul>

          <li class="main-menu-item">
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Learn
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/get-started">
                  <span class=dropdown-title>Get Started</span>
                  <p>Run PyTorch locally or get started quickly with one of the supported cloud platforms</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials">
                  <span class="dropdown-title">Tutorials</span>
                  <p>Whats new in PyTorch tutorials</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials/beginner/basics/intro.html">
                  <span class="dropdown-title">Learn the Basics</span>
                  <p>Familiarize yourself with PyTorch concepts and modules</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials/recipes/recipes_index.html">
                  <span class="dropdown-title">PyTorch Recipes</span>
                  <p>Bite-size, ready-to-deploy PyTorch code examples</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/tutorials/beginner/introyt.html">
                  <span class="dropdown-title">Intro to PyTorch - YouTube Series</span>
                  <p>Master PyTorch basics with our engaging YouTube tutorial series</p>
                </a>
              </div>
            </div>
          </li>

          <li>
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Ecosystem
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem">
                  <span class="dropdown-title">Tools</span>
                  <p>Learn about the tools and frameworks in the PyTorch Ecosystem</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/#community-module">
                  <span class=dropdown-title>Community</span>
                  <p>Join the PyTorch developer community to contribute, learn, and get your questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://discuss.pytorch.org/" target="_blank">
                  <span class=dropdown-title>Forums</span>
                  <p>A place to discuss PyTorch code, issues, install, research</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/resources">
                  <span class=dropdown-title>Developer Resources</span>
                  <p>Find resources and get questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem/contributor-awards-2023">
                  <span class="dropdown-title">Contributor Awards - 2023</span>
                  <p>Award winners announced at this year's PyTorch Conference</p>
                </a>
              </div>
            </div>
          </li>

          <li>
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Edge
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/edge">
                  <span class="dropdown-title">About PyTorch Edge</span>
                  <p>Build innovative and privacy-aware AI experiences for edge devices</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/executorch-overview">
                  <span class="dropdown-title">ExecuTorch</span>
                  <p>End-to-end solution for enabling on-device inference capabilities across mobile and edge devices</p>
                </a>
              </div>
            </div>  
          </li>

          <li class="main-menu-item">
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Docs
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/docs/stable/index.html">
                  <span class="dropdown-title">PyTorch</span>
                  <p>Explore the documentation for comprehensive guidance on how to use PyTorch</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/pytorch-domains">
                  <span class="dropdown-title">PyTorch Domains</span>
                  <p>Read the PyTorch Domains documentation to learn more about domain-specific libraries</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                Blogs & News 
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/blog/">
                  <span class="dropdown-title">PyTorch Blog</span>
                  <p>Catch up on the latest technical news and happenings</p>
                </a>
                 <a class="nav-dropdown-item" href="https://pytorch.org/community-blog">
                  <span class="dropdown-title">Community Blog</span>
                  <p>Stories from the PyTorch ecosystem</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/videos">
                  <span class="dropdown-title">Videos</span>
                  <p>Learn about the latest PyTorch tutorials, new, and more </p>
                <a class="nav-dropdown-item" href="https://pytorch.org/community-stories">
                  <span class="dropdown-title">Community Stories</span>
                  <p>Learn how our community solves real, everyday machine learning problems with PyTorch</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/events">
                  <span class="dropdown-title">Events</span>
                  <p>Find events, webinars, and podcasts</p>
                </a>
            </div>
          </li>

          <li>
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="with-down-arrow">
                About
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/foundation">
                  <span class="dropdown-title">PyTorch Foundation</span>
                  <p>Learn more about the PyTorch Foundation</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/governing-board">
                  <span class="dropdown-title">Governing Board</span>
                  <p></p>
                </a>
              </div>
            </div>
          </li>

          <li class="main-menu-item">
            <div class="no-dropdown">
              <a href="https://pytorch.org/join" data-cta="join">
                Become a Member
              </a>
            </div>
          </li>
          <li>
           <div class="main-menu-item">
             <a href="https://github.com/pytorch/pytorch" class="github-icon">
             </a>
           </div>
          </li>
          <!--- TODO: This block adds the search icon to the nav bar. We will enable it later. 
          <li>
            <div class="main-menu-item">
             <a href="https://github.com/pytorch/pytorch" class="search-icon">
             </a>
            </div>
          </li>
          --->
        </ul>
      </div>

      <a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
    </div>
  </div>
</div>

<body class="pytorch-body">

   
    <div class="table-of-contents-link-wrapper">
      <span>Table of Contents</span>
      <a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
    </div>

    <nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
      <div class="pytorch-side-scroll">
        <div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          <div class="pytorch-left-menu-search">
            
    <div class="version">
      <a href='https://pytorch.org/docs/versions.html'>2.6 &#x25BC</a>
    </div>
    <div id="searchBox">
    <div class="searchbox" id="googleSearchBox">
      <script async src="https://cse.google.com/cse.js?cx=e65585f8c3ea1440e"></script>
      <div class="gcse-search"></div>
    </div>
    <div id="sphinxSearchBox" style="display: none;">
      <div role="search">
        <form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
          <input type="text" name="q" placeholder="Search Docs" />
          <input type="hidden" name="check_keywords" value="yes" />
          <input type="hidden" name="area" value="default" />
        </form>
      </div>
    </div>
  </div>
  <form id="searchForm">
    <label style="margin-bottom: 1rem">
      <input type="radio" name="searchType" value="google" checked>
      Google Search
    </label>
    <label style="margin-bottom: 1rem">
      <input type="radio" name="searchType" value="sphinx">
      Classic Search
    </label>
  </form>

  <script>
     document.addEventListener('DOMContentLoaded', function() {
      const searchForm = document.getElementById('searchForm');
      const googleSearchBox = document.getElementById('googleSearchBox');
      const sphinxSearchBox = document.getElementById('sphinxSearchBox');
      // Function to toggle search box visibility
      function toggleSearchBox(searchType) {
        googleSearchBox.style.display = searchType === 'google' ? 'block' : 'none';
        sphinxSearchBox.style.display = searchType === 'sphinx' ? 'block' : 'none';
      }
      // Determine the default search type
      let defaultSearchType;
      const currentUrl = window.location.href;
      if (currentUrl.startsWith('https://pytorch.org/docs/stable')) {
        // For the stable documentation, default to Google
        defaultSearchType = localStorage.getItem('searchType') || 'google';
      } else {
        // For any other version, including docs-preview, default to Sphinx
        defaultSearchType = 'sphinx';
      }
      // Set the default search type
      document.querySelector(`input[name="searchType"][value="${defaultSearchType}"]`).checked = true;
      toggleSearchBox(defaultSearchType);
      // Event listener for changes in search type
      searchForm.addEventListener('change', function(event) {
        const selectedSearchType = event.target.value;
        localStorage.setItem('searchType', selectedSearchType);
        toggleSearchBox(selectedSearchType);
      });
      // Set placeholder text for Google search box
      window.onload = function() {
        var placeholderText = "Search Docs";
        var googleSearchboxText = document.querySelector("#gsc-i-id1");
        if (googleSearchboxText) {
          googleSearchboxText.placeholder = placeholderText;
          googleSearchboxText.style.fontFamily = 'FreightSans';
          googleSearchboxText.style.fontSize = "1.2rem";
          googleSearchboxText.style.color = '#262626';
        }
      };
    });
  </script>
     <!--temporarily add a link to survey -->
     <div class="survey-link">
       <p><i class="fas fa-poll" aria-hidden="true">&nbsp </i><a href="https://forms.gle/tdrnwJhaQ9tUePxz9">Share Your Feedback</a> about our new search</p>
     </div>

          </div>

          
              <p class="caption" role="heading"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../community/build_ci_governance.html">PyTorch Governance | Build + CI</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/design.html">PyTorch Design Philosophy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/governance.html">PyTorch Governance | Mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/persons_of_interest.html">PyTorch Governance | Maintainers</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Developer Notes</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/amp_examples.html">Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/custom_operators.html">PyTorch Custom Operators Landing Page</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/extending.func.html">Extending torch.func with autograd.Function</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/fsdp.html">FSDP Notes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/get_start_xpu.html">Getting Started on Intel GPU</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/gradcheck.html">Gradcheck mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/hip.html">HIP (ROCm) semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/modules.html">Modules</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/mps.html">MPS backend</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/numerical_accuracy.html">Numerical accuracy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/windows.html">Windows FAQ</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../cpp_index.html">C++</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/javadoc/">Javadoc</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../deploy.html">torch::deploy</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../amp.html">torch.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../library.html">torch.library</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../accelerator.html">torch.accelerator</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cpu.html">torch.cpu</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch_cuda_memory.html">Understanding CUDA Memory Usage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch_cuda_memory.html#generating-a-snapshot">Generating a Snapshot</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch_cuda_memory.html#using-the-visualizer">Using the visualizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch_cuda_memory.html#snapshot-api-reference">Snapshot API Reference</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../mps.html">torch.mps</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../xpu.html">torch.xpu</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../mtia.html">torch.mtia</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../mtia.memory.html">torch.mtia.memory</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../meta.html">Meta device</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../backends.html">torch.backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../export.html">torch.export</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.tensor.html">torch.distributed.tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.algorithms.join.html">torch.distributed.algorithms.join</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.elastic.html">torch.distributed.elastic</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../fsdp.html">torch.distributed.fsdp</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.fsdp.fully_shard.html">torch.distributed.fsdp.fully_shard</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.tensor.parallel.html">torch.distributed.tensor.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.optim.html">torch.distributed.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.pipelining.html">torch.distributed.pipelining</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.checkpoint.html">torch.distributed.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch.compiler.html">torch.compiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../fft.html">torch.fft</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../func.html">torch.func</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../futures.html">torch.futures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../fx.html">torch.fx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../fx.experimental.html">torch.fx.experimental</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../hub.html">torch.hub</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../linalg.html">torch.linalg</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../monitor.html">torch.monitor</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../signal.html">torch.signal</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../special.html">torch.special</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch.overrides.html">torch.overrides</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../package.html">torch.package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../profiler.html">torch.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.attention.html">torch.nn.attention</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../complex_numbers.html">Complex Numbers</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../ddp_comm_hooks.html">DDP Communication Hooks</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../rpc.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../masked.html">torch.masked</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nested.html">torch.nested</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../size.html">torch.Size</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../testing.html">torch.testing</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../utils.html">torch.utils</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../benchmark_utils.html">torch.utils.benchmark</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../deterministic.html">torch.utils.deterministic</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../jit_utils.html">torch.utils.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../mobile_optimizer.html">torch.utils.mobile_optimizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../module_tracker.html">torch.utils.module_tracker</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../config_mod.html">torch.__config__</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../future_mod.html">torch.__future__</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../logging.html">torch._logging</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../torch_environment_variables.html">Torch Environment Variables</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio/stable">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/data">TorchData</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/torchrec">TorchRec</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text/stable">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/vision/stable">torchvision</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>

            
        </div>
      </div>
    </nav>

    <div class="pytorch-container">
      <div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
        <div class="pytorch-breadcrumbs-wrapper">
          

<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="pytorch-breadcrumbs">
    
      <li>
        <a href="../../../index.html">
          
            Docs
          
        </a> &gt;
      </li>

        
          <li><a href="../../index.html">Module code</a> &gt;</li>
        
          <li><a href="../../torch.html">torch</a> &gt;</li>
        
          <li><a href="../autograd.html">torch.autograd</a> &gt;</li>
        
      <li>torch.autograd.graph</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
</div>
        </div>

        <div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
          Shortcuts
        </div>
      </div>

      <section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
        <div class="pytorch-content-left">

        
          <!-- Google Tag Manager (noscript) -->
          <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-T8XT4PS"
          height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>
          <!-- End Google Tag Manager (noscript) -->
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <h1>Source code for torch.autograd.graph</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">abc</span>
<span class="kn">import</span> <span class="nn">contextlib</span>
<span class="kn">import</span> <span class="nn">functools</span>
<span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">threading</span>
<span class="kn">from</span> <span class="nn">collections</span> <span class="kn">import</span> <span class="n">defaultdict</span><span class="p">,</span> <span class="n">deque</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">Any</span><span class="p">,</span>
    <span class="n">Callable</span><span class="p">,</span>
    <span class="n">cast</span><span class="p">,</span>
    <span class="n">Deque</span><span class="p">,</span>
    <span class="n">Dict</span><span class="p">,</span>
    <span class="n">Generator</span><span class="p">,</span>
    <span class="n">Iterable</span><span class="p">,</span>
    <span class="n">Iterator</span><span class="p">,</span>
    <span class="n">List</span><span class="p">,</span>
    <span class="n">Literal</span><span class="p">,</span>
    <span class="n">MutableMapping</span><span class="p">,</span>
    <span class="n">NamedTuple</span><span class="p">,</span>
    <span class="n">Optional</span><span class="p">,</span>
    <span class="n">Sequence</span><span class="p">,</span>
    <span class="n">Set</span><span class="p">,</span>
    <span class="n">Tuple</span><span class="p">,</span>
    <span class="n">TYPE_CHECKING</span><span class="p">,</span>
    <span class="n">Union</span><span class="p">,</span>
<span class="p">)</span>
<span class="kn">from</span> <span class="nn">typing_extensions</span> <span class="kn">import</span> <span class="n">TypeAlias</span>
<span class="kn">from</span> <span class="nn">weakref</span> <span class="kn">import</span> <span class="n">WeakKeyDictionary</span><span class="p">,</span> <span class="n">WeakValueDictionary</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.autograd.variable</span> <span class="kn">import</span> <span class="n">Variable</span>
<span class="kn">from</span> <span class="nn">torch.utils._python_dispatch</span> <span class="kn">import</span> <span class="n">TorchDispatchMode</span>
<span class="kn">from</span> <span class="nn">torch.utils.hooks</span> <span class="kn">import</span> <span class="n">RemovableHandle</span>


<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="kn">from</span> <span class="nn">torch._ops</span> <span class="kn">import</span> <span class="n">OpOverload</span>


<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s2">&quot;saved_tensors_hooks&quot;</span><span class="p">,</span>
    <span class="s2">&quot;save_on_cpu&quot;</span><span class="p">,</span>
    <span class="s2">&quot;disable_saved_tensors_hooks&quot;</span><span class="p">,</span>
    <span class="s2">&quot;register_multi_grad_hook&quot;</span><span class="p">,</span>
    <span class="s2">&quot;allow_mutation_on_saved_tensors&quot;</span><span class="p">,</span>
    <span class="s2">&quot;Node&quot;</span><span class="p">,</span>
    <span class="s2">&quot;GradientEdge&quot;</span><span class="p">,</span>
    <span class="s2">&quot;get_gradient_edge&quot;</span><span class="p">,</span>
    <span class="s2">&quot;increment_version&quot;</span><span class="p">,</span>
<span class="p">]</span>


<span class="n">log</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">Node</span><span class="p">(</span><span class="n">abc</span><span class="o">.</span><span class="n">ABC</span><span class="p">):</span>
<div class="viewcode-block" id="Node.name"><a class="viewcode-back" href="../../../generated/torch.autograd.graph.Node.name.html#torch.autograd.graph.Node.name">[docs]</a>    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">name</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
<span class="w">        </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Return the name.</span>

<span class="sd">        Example::</span>

<span class="sd">            &gt;&gt;&gt; import torch</span>
<span class="sd">            &gt;&gt;&gt; a = torch.tensor([0., 0., 0.], requires_grad=True)</span>
<span class="sd">            &gt;&gt;&gt; b = a.clone()</span>
<span class="sd">            &gt;&gt;&gt; assert isinstance(b.grad_fn, torch.autograd.graph.Node)</span>
<span class="sd">            &gt;&gt;&gt; print(b.grad_fn.name())</span>
<span class="sd">            CloneBackward0</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

    <span class="nd">@property</span>
    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">next_functions</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">Optional</span><span class="p">[</span><span class="s2">&quot;Node&quot;</span><span class="p">],</span> <span class="nb">int</span><span class="p">],</span> <span class="o">...</span><span class="p">]:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>

<div class="viewcode-block" id="Node.metadata"><a class="viewcode-back" href="../../../generated/torch.autograd.graph.Node.metadata.html#torch.autograd.graph.Node.metadata">[docs]</a>    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">metadata</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">:</span>
<span class="w">        </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Return the metadata.&quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

    <span class="nd">@property</span>
    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">_input_metadata</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">]:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>

    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">_register_hook_dict</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span>

<div class="viewcode-block" id="Node.register_hook"><a class="viewcode-back" href="../../../generated/torch.autograd.graph.Node.register_hook.html#torch.autograd.graph.Node.register_hook">[docs]</a>    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">register_hook</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fn</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">RemovableHandle</span><span class="p">:</span>
<span class="w">        </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Register a backward hook.</span>

<span class="sd">        The hook will be called every time a gradient with respect to the</span>
<span class="sd">        Node is computed. The hook should have the following signature::</span>

<span class="sd">            hook(grad_inputs: Tuple[Tensor], grad_outputs: Tuple[Tensor]) -&gt; Tuple[Tensor] or None</span>


<span class="sd">        The hook should not modify its argument, but it can optionally return</span>
<span class="sd">        a new gradient which will be used in place of :attr:`grad_inputs`.</span>

<span class="sd">        This function returns a handle with a method ``handle.remove()``</span>
<span class="sd">        that removes the hook from the module.</span>

<span class="sd">        .. note::</span>
<span class="sd">            See :ref:`backward-hooks-execution` for more information on how when this hook</span>
<span class="sd">            is executed, and how its execution is ordered relative to other hooks.</span>

<span class="sd">        .. note::</span>
<span class="sd">            In the rare case where the hook is registered while the Node has already</span>
<span class="sd">            begun execution, there is no longer any guarantee on :attr:`grad_outputs`</span>
<span class="sd">            content (it might be as usual or empty depending on other factors). The</span>
<span class="sd">            hook can still optionally return a new gradient to be used in place of</span>
<span class="sd">            :attr:`grad_inputs` independent of :attr:`grad_outputs`.</span>

<span class="sd">        Example::</span>

<span class="sd">            &gt;&gt;&gt; import torch</span>
<span class="sd">            &gt;&gt;&gt; a = torch.tensor([0., 0., 0.], requires_grad=True)</span>
<span class="sd">            &gt;&gt;&gt; b = a.clone()</span>
<span class="sd">            &gt;&gt;&gt; assert isinstance(b.grad_fn, torch.autograd.graph.Node)</span>
<span class="sd">            &gt;&gt;&gt; handle = b.grad_fn.register_hook(lambda gI, gO: (gO[0] * 2,))</span>
<span class="sd">            &gt;&gt;&gt; b.sum().backward(retain_graph=True)</span>
<span class="sd">            &gt;&gt;&gt; print(a.grad)</span>
<span class="sd">            tensor([2., 2., 2.])</span>
<span class="sd">            &gt;&gt;&gt; handle.remove() # Removes the hook</span>
<span class="sd">            &gt;&gt;&gt; a.grad = None</span>
<span class="sd">            &gt;&gt;&gt; b.sum().backward(retain_graph=True)</span>
<span class="sd">            &gt;&gt;&gt; print(a.grad)</span>
<span class="sd">            tensor([1., 1., 1.])</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

<div class="viewcode-block" id="Node.register_prehook"><a class="viewcode-back" href="../../../generated/torch.autograd.graph.Node.register_prehook.html#torch.autograd.graph.Node.register_prehook">[docs]</a>    <span class="nd">@abc</span><span class="o">.</span><span class="n">abstractmethod</span>
    <span class="k">def</span> <span class="nf">register_prehook</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fn</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[</span><span class="o">...</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">RemovableHandle</span><span class="p">:</span>
<span class="w">        </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Register a backward pre-hook.</span>

<span class="sd">        The hook will be called every time a gradient with respect to the</span>
<span class="sd">        Node is computed. The hook should have the following signature::</span>

<span class="sd">            hook(grad_outputs: Tuple[Tensor]) -&gt; Tuple[Tensor] or None</span>

<span class="sd">        The hook should not modify its argument, but it can optionally return</span>
<span class="sd">        a new gradient which will be used in place of :attr:`grad_outputs`.</span>

<span class="sd">        This function returns a handle with a method ``handle.remove()``</span>
<span class="sd">        that removes the hook from the module.</span>

<span class="sd">        .. note::</span>
<span class="sd">            See :ref:`backward-hooks-execution` for more information on how when this hook</span>
<span class="sd">            is executed, and how its execution is ordered relative to other hooks.</span>

<span class="sd">        Example::</span>

<span class="sd">            &gt;&gt;&gt; a = torch.tensor([0., 0., 0.], requires_grad=True)</span>
<span class="sd">            &gt;&gt;&gt; b = a.clone()</span>
<span class="sd">            &gt;&gt;&gt; assert isinstance(b.grad_fn, torch.autograd.graph.Node)</span>
<span class="sd">            &gt;&gt;&gt; handle = b.grad_fn.register_prehook(lambda gI: (gI[0] * 2,))</span>
<span class="sd">            &gt;&gt;&gt; b.sum().backward(retain_graph=True)</span>
<span class="sd">            &gt;&gt;&gt; print(a.grad)</span>
<span class="sd">            tensor([2., 2., 2.])</span>
<span class="sd">            &gt;&gt;&gt; handle.remove()</span>
<span class="sd">            &gt;&gt;&gt; a.grad = None</span>
<span class="sd">            &gt;&gt;&gt; b.sum().backward(retain_graph=True)</span>
<span class="sd">            &gt;&gt;&gt; print(a.grad)</span>
<span class="sd">            tensor([1., 1., 1.])</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span></div>

    <span class="nd">@classmethod</span>
    <span class="k">def</span> <span class="nf">__subclasshook__</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">subclass</span><span class="p">:</span> <span class="nb">type</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
        <span class="k">if</span> <span class="bp">cls</span> <span class="ow">is</span> <span class="n">Node</span> <span class="ow">and</span> <span class="p">(</span>
            <span class="p">(</span>
                <span class="n">subclass</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                <span class="ow">and</span> <span class="n">subclass</span> <span class="ow">is</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functions</span><span class="p">,</span> <span class="n">subclass</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
            <span class="p">)</span>
            <span class="ow">or</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">subclass</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">autograd</span><span class="o">.</span><span class="n">function</span><span class="o">.</span><span class="n">BackwardCFunction</span><span class="p">)</span>
        <span class="p">):</span>
            <span class="k">return</span> <span class="kc">True</span>
        <span class="k">return</span> <span class="bp">NotImplemented</span>


<span class="k">def</span> <span class="nf">_get_grad_fn_or_grad_acc</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="s2">&quot;GradientEdge&quot;</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Node</span><span class="p">:</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">GradientEdge</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">t</span><span class="o">.</span><span class="n">node</span>
    <span class="k">if</span> <span class="n">t</span><span class="o">.</span><span class="n">requires_grad</span> <span class="ow">and</span> <span class="n">t</span><span class="o">.</span><span class="n">grad_fn</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">enable_grad</span><span class="p">():</span>
            <span class="n">node</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">view_as</span><span class="p">(</span><span class="n">t</span><span class="p">)</span><span class="o">.</span><span class="n">grad_fn</span><span class="o">.</span><span class="n">next_functions</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>  <span class="c1"># type: ignore[union-attr]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">node</span> <span class="o">=</span> <span class="n">t</span><span class="o">.</span><span class="n">grad_fn</span>
    <span class="k">assert</span> <span class="n">node</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
    <span class="k">return</span> <span class="n">node</span>


<div class="viewcode-block" id="GradientEdge"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.GradientEdge">[docs]</a><span class="k">class</span> <span class="nc">GradientEdge</span><span class="p">(</span><span class="n">NamedTuple</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Object representing a given gradient edge within the autograd graph.</span>

<span class="sd">    To get the gradient edge where a given Tensor gradient will be computed,</span>
<span class="sd">    you can do ``edge = autograd.graph.get_gradient_edge(tensor)``.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">node</span><span class="p">:</span> <span class="n">Node</span>
    <span class="n">output_nr</span><span class="p">:</span> <span class="nb">int</span></div>


<div class="viewcode-block" id="get_gradient_edge"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.get_gradient_edge">[docs]</a><span class="k">def</span> <span class="nf">get_gradient_edge</span><span class="p">(</span><span class="n">tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">GradientEdge</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Get the gradient edge for computing the gradient of the given Tensor.</span>

<span class="sd">    In particular, it is equivalent to call</span>
<span class="sd">    ``g = autograd.grad(loss, input)`` and ``g = autograd.grad(loss, get_gradient_edge(input))``.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">tensor</span><span class="o">.</span><span class="n">requires_grad</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
            <span class="s2">&quot;It is not possible to get the gradient edge for a Tensor &quot;</span>
            <span class="s2">&quot;that does not require gradients&quot;</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="n">grad_fn</span> <span class="o">=</span> <span class="n">_get_grad_fn_or_grad_acc</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>

    <span class="c1"># Note that output_nr default to 0 which is the right value</span>
    <span class="c1"># for the AccumulateGrad node.</span>
    <span class="k">return</span> <span class="n">GradientEdge</span><span class="p">(</span><span class="n">grad_fn</span><span class="p">,</span> <span class="n">tensor</span><span class="o">.</span><span class="n">output_nr</span><span class="p">)</span></div>


<div class="viewcode-block" id="increment_version"><a class="viewcode-back" href="../../../generated/torch.autograd.graph.increment_version.html#torch.autograd.graph.increment_version">[docs]</a><span class="k">def</span> <span class="nf">increment_version</span><span class="p">(</span><span class="n">tensor</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Iterable</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Update autograd metadata tracking whether the given Tensor was modified in place.</span>

<span class="sd">    This is to enable more accurate error checking within the autograd engine.</span>
<span class="sd">    It is already done automatically by PyTorch functions and within custom Function</span>
<span class="sd">    when mark_dirty() is called appropriately so you only need to call this explicitly</span>
<span class="sd">    if you are doing inplace operation on the Tensor data in a way that Pytorch doesn&#39;t</span>
<span class="sd">    know about. For example a custom kernel that reads the Tensor data_ptr and modifies</span>
<span class="sd">    the memory inplace based on this pointer. Can accept either a tensor, or a list of tensors.</span>

<span class="sd">    Note that incrementing the version counter multiple times for a single inplace operation</span>
<span class="sd">    is not problematic.</span>

<span class="sd">    Note that if you pass in tensor constructed under torch.inference_mode(),</span>
<span class="sd">    we will not bump its version counter (because your tensor does not have one).</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="n">tensor</span> <span class="o">=</span> <span class="p">(</span><span class="n">tensor</span><span class="p">,)</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_increment_version</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span></div>


<div class="viewcode-block" id="saved_tensors_hooks"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.saved_tensors_hooks">[docs]</a><span class="k">class</span> <span class="nc">saved_tensors_hooks</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Context-manager that sets a pair of pack / unpack hooks for saved tensors.</span>

<span class="sd">    Use this context-manager to define how intermediary results of an operation</span>
<span class="sd">    should be packed before saving, and unpacked on retrieval.</span>

<span class="sd">    In that context, the ``pack_hook`` function will be called everytime an</span>
<span class="sd">    operation saves a tensor for backward (this includes intermediary results</span>
<span class="sd">    saved using</span>
<span class="sd">    :func:`~torch.autograd.function._ContextMethodMixin.save_for_backward` but</span>
<span class="sd">    also those recorded by a PyTorch-defined operation). The output of</span>
<span class="sd">    ``pack_hook`` is then stored in the computation graph instead of the</span>
<span class="sd">    original tensor.</span>

<span class="sd">    The ``unpack_hook`` is called when the saved tensor needs to be accessed,</span>
<span class="sd">    namely when executing :func:`torch.Tensor.backward()` or</span>
<span class="sd">    :func:`torch.autograd.grad()`. It takes as argument the *packed* object</span>
<span class="sd">    returned by ``pack_hook`` and should return a tensor which has the same</span>
<span class="sd">    content as the original tensor (passed as input to the corresponding</span>
<span class="sd">    ``pack_hook``).</span>

<span class="sd">    The hooks should have the following signatures:</span>

<span class="sd">        pack_hook(tensor: Tensor) -&gt; Any</span>

<span class="sd">        unpack_hook(Any) -&gt; Tensor</span>

<span class="sd">    where the return value of ``pack_hook`` is a valid input to ``unpack_hook``.</span>

<span class="sd">    In general, you want ``unpack_hook(pack_hook(t))`` to be equal to ``t`` in terms</span>
<span class="sd">    of value, size, dtype and device.</span>

<span class="sd">    Example::</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)</span>
<span class="sd">        &gt;&gt;&gt; def pack_hook(x):</span>
<span class="sd">        ...     print(&quot;Packing&quot;, x)</span>
<span class="sd">        ...     return x</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; def unpack_hook(x):</span>
<span class="sd">        ...     print(&quot;Unpacking&quot;, x)</span>
<span class="sd">        ...     return x</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; a = torch.ones(5, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.ones(5, requires_grad=True) * 2</span>
<span class="sd">        &gt;&gt;&gt; with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):</span>
<span class="sd">        ...     y = a * b</span>
<span class="sd">        Packing tensor([1., 1., 1., 1., 1.], requires_grad=True)</span>
<span class="sd">        Packing tensor([2., 2., 2., 2., 2.], grad_fn=&lt;MulBackward0&gt;)</span>
<span class="sd">        &gt;&gt;&gt; y.sum().backward()</span>
<span class="sd">        Unpacking tensor([1., 1., 1., 1., 1.], requires_grad=True)</span>
<span class="sd">        Unpacking tensor([2., 2., 2., 2., 2.], grad_fn=&lt;MulBackward0&gt;)</span>

<span class="sd">    .. warning ::</span>
<span class="sd">        Performing an inplace operation on the input to either hooks may lead</span>
<span class="sd">        to undefined behavior.</span>

<span class="sd">    .. warning ::</span>
<span class="sd">        Only one pair of hooks is allowed at a time. When recursively nesting this</span>
<span class="sd">        context-manager, only the inner-most pair of hooks will be applied.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">pack_hook</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">Any</span><span class="p">],</span>
        <span class="n">unpack_hook</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[[</span><span class="n">Any</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pack_hook</span> <span class="o">=</span> <span class="n">pack_hook</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">unpack_hook</span> <span class="o">=</span> <span class="n">unpack_hook</span>

    <span class="k">def</span> <span class="fm">__enter__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_autograd</span><span class="o">.</span><span class="n">_push_saved_tensors_default_hooks</span><span class="p">(</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">pack_hook</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">unpack_hook</span>
        <span class="p">)</span>

    <span class="k">def</span> <span class="fm">__exit__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="nb">object</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_autograd</span><span class="o">.</span><span class="n">_pop_saved_tensors_default_hooks</span><span class="p">()</span></div>


<div class="viewcode-block" id="save_on_cpu"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.save_on_cpu">[docs]</a><span class="k">class</span> <span class="nc">save_on_cpu</span><span class="p">(</span><span class="n">saved_tensors_hooks</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Context manager under which tensors saved by the forward pass will be stored on cpu, then retrieved for backward.</span>

<span class="sd">    When performing operations within this context manager, intermediary</span>
<span class="sd">    results saved in the graph during the forward pass will be moved to CPU,</span>
<span class="sd">    then copied back to the original device when needed for the backward pass.</span>
<span class="sd">    If the graph was already on CPU, no tensor copy is performed.</span>

<span class="sd">    Use this context-manager to trade compute for GPU memory usage (e.g.</span>
<span class="sd">    when your model doesn&#39;t fit in GPU memory during training).</span>

<span class="sd">    Args:</span>
<span class="sd">        pin_memory (bool): If ``True`` tensors will be saved to CPU pinned memory</span>
<span class="sd">                           during packing and copied to GPU asynchronously during unpacking.</span>
<span class="sd">                           Defaults to ``False``.</span>
<span class="sd">                           Also see :ref:`cuda-memory-pinning`.</span>


<span class="sd">    Example::</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)</span>
<span class="sd">        &gt;&gt;&gt; # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)</span>
<span class="sd">        &gt;&gt;&gt; a = torch.randn(5, requires_grad=True, device=&quot;cuda&quot;)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.randn(5, requires_grad=True, device=&quot;cuda&quot;)</span>
<span class="sd">        &gt;&gt;&gt; c = torch.randn(5, requires_grad=True, device=&quot;cuda&quot;)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; def f(a, b, c):</span>
<span class="sd">        ...     prod_1 = a * b           # a and b are saved on GPU</span>
<span class="sd">        ...     with torch.autograd.graph.save_on_cpu():</span>
<span class="sd">        ...         prod_2 = prod_1 * c  # prod_1 and c are saved on CPU</span>
<span class="sd">        ...     y = prod_2 * a           # prod_2 and a are saved on GPU</span>
<span class="sd">        ...     return y</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; y = f(a, b, c)</span>
<span class="sd">        &gt;&gt;&gt; del a, b, c  # for illustration only</span>
<span class="sd">        &gt;&gt;&gt; # the content of a, b, and prod_2 are still alive on GPU</span>
<span class="sd">        &gt;&gt;&gt; # the content of prod_1 and c only live on CPU</span>
<span class="sd">        &gt;&gt;&gt; y.sum().backward()  # all CPU tensors are moved back to GPU, for backward</span>
<span class="sd">        &gt;&gt;&gt; # all intermediary tensors are released (deleted) after the call to backward</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">pin_memory</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> <span class="n">device_type</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;cuda&quot;</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">device_module</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">torch</span><span class="p">,</span> <span class="n">device_type</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">cuda</span><span class="p">)</span>

        <span class="k">def</span> <span class="nf">pack_to_cpu</span><span class="p">(</span><span class="n">tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">pin_memory</span><span class="p">:</span>
                <span class="k">return</span> <span class="p">(</span><span class="n">tensor</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">tensor</span><span class="o">.</span><span class="n">cpu</span><span class="p">())</span>
            <span class="n">packed</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">empty</span><span class="p">(</span>
                <span class="n">tensor</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span>
                <span class="n">dtype</span><span class="o">=</span><span class="n">tensor</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span>
                <span class="n">layout</span><span class="o">=</span><span class="n">tensor</span><span class="o">.</span><span class="n">layout</span><span class="p">,</span>
                <span class="n">pin_memory</span><span class="o">=</span><span class="p">(</span><span class="n">device_module</span><span class="o">.</span><span class="n">is_available</span><span class="p">()</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">tensor</span><span class="o">.</span><span class="n">is_sparse</span><span class="p">),</span>
            <span class="p">)</span>
            <span class="n">packed</span><span class="o">.</span><span class="n">copy_</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">tensor</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">packed</span><span class="p">)</span>

        <span class="k">def</span> <span class="nf">unpack_from_cpu</span><span class="p">(</span><span class="n">packed</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
            <span class="n">device</span><span class="p">,</span> <span class="n">tensor</span> <span class="o">=</span> <span class="n">packed</span>
            <span class="k">return</span> <span class="n">tensor</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">device</span><span class="p">,</span> <span class="n">non_blocking</span><span class="o">=</span><span class="n">pin_memory</span><span class="p">)</span>

        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">pack_to_cpu</span><span class="p">,</span> <span class="n">unpack_from_cpu</span><span class="p">)</span></div>


<div class="viewcode-block" id="disable_saved_tensors_hooks"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.disable_saved_tensors_hooks">[docs]</a><span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
<span class="k">def</span> <span class="nf">disable_saved_tensors_hooks</span><span class="p">(</span><span class="n">error_message</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Generator</span><span class="p">[</span><span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Context-manager that disables the saved tensors default hooks feature.</span>

<span class="sd">    Useful for if you are creating a feature that does not work with saved</span>
<span class="sd">    tensors default hooks.</span>

<span class="sd">    Args:</span>
<span class="sd">        error_message (str): When saved tensors default hooks are used when they</span>
<span class="sd">                             have been are disabled, a RuntimeError with this</span>
<span class="sd">                             error message gets raised.</span>

<span class="sd">    Example::</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +SKIP(failing)</span>
<span class="sd">        &gt;&gt;&gt; message = &quot;saved tensors default hooks are disabled&quot;</span>
<span class="sd">        &gt;&gt;&gt; with torch.autograd.graph.disable_saved_tensors_hooks(message):</span>
<span class="sd">        ...     # Raises RuntimeError: saved tensors default hooks are disabled</span>
<span class="sd">        ...     with torch.autograd.graph.save_on_cpu():</span>
<span class="sd">        ...         pass</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">maybe_prev_message</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">maybe_prev_message</span> <span class="o">=</span> <span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_autograd</span><span class="o">.</span><span class="n">_saved_tensors_hooks_get_disabled_error_message</span><span class="p">()</span>
        <span class="p">)</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_autograd</span><span class="o">.</span><span class="n">_saved_tensors_hooks_disable</span><span class="p">(</span><span class="n">error_message</span><span class="p">)</span>
        <span class="k">yield</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="c1"># See NOTE: [disabled_error_message invariant]</span>
        <span class="k">if</span> <span class="n">maybe_prev_message</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_autograd</span><span class="o">.</span><span class="n">_saved_tensors_hooks_enable</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_autograd</span><span class="o">.</span><span class="n">_saved_tensors_hooks_disable</span><span class="p">(</span><span class="n">maybe_prev_message</span><span class="p">)</span></div>


<span class="k">class</span> <span class="nc">_MultiHandle</span><span class="p">(</span><span class="n">RemovableHandle</span><span class="p">):</span>
    <span class="n">handles</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">RemovableHandle</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">handles</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">RemovableHandle</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">handles</span> <span class="o">=</span> <span class="n">handles</span>

    <span class="k">def</span> <span class="nf">remove</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">handle</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">handles</span><span class="p">:</span>
            <span class="n">handle</span><span class="o">.</span><span class="n">remove</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">__getstate__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">RemovableHandle</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">handles</span>

    <span class="k">def</span> <span class="nf">__setstate__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">state</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">RemovableHandle</span><span class="p">,</span> <span class="o">...</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">handles</span> <span class="o">=</span> <span class="n">state</span>


<div class="viewcode-block" id="register_multi_grad_hook"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.register_multi_grad_hook">[docs]</a><span class="k">def</span> <span class="nf">register_multi_grad_hook</span><span class="p">(</span>
    <span class="n">tensors</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">fn</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span>
        <span class="n">Callable</span><span class="p">[[</span><span class="n">Sequence</span><span class="p">[</span><span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]]],</span> <span class="kc">None</span><span class="p">],</span>
        <span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="kc">None</span><span class="p">],</span>
    <span class="p">],</span>
    <span class="o">*</span><span class="p">,</span>
    <span class="n">mode</span><span class="p">:</span> <span class="n">Literal</span><span class="p">[</span><span class="s2">&quot;all&quot;</span><span class="p">,</span> <span class="s2">&quot;any&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;all&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">RemovableHandle</span><span class="p">:</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Register a multi-grad backward hook.</span>

<span class="sd">    There are two supported modes: ``&quot;all&quot;`` and ``&quot;any&quot;``.</span>

<span class="sd">    Under the ``&quot;all&quot;`` mode, the hook will be called after gradients with respect to every tensor in</span>
<span class="sd">    :attr:`tensors` have been computed. If a tensor is in :attr:`tensors` but</span>
<span class="sd">    is not part of the graph, or if a tensor is not needed to compute the gradients</span>
<span class="sd">    for any ``inputs`` specified for the current ``.backward()`` or ``.grad()`` call,</span>
<span class="sd">    this tensor will be ignored and the hook will not wait for its gradient to be</span>
<span class="sd">    computed.</span>

<span class="sd">    After every non-ignored tensor&#39;s gradient has been computed, :attr:`fn` will be</span>
<span class="sd">    called with those gradients. ``None`` will be passed for tensors that did not</span>
<span class="sd">    have their gradients computed.</span>

<span class="sd">    Under the ``&quot;any&quot;`` mode, the hook will be called after the first gradient</span>
<span class="sd">    with respect to a tensor in :attr:`tensors` has been computed. The hook</span>
<span class="sd">    will be called with that gradient as its argument.</span>

<span class="sd">    The hook should not modify its arguments.</span>

<span class="sd">    This function returns a handle with a method ``handle.remove()`` that removes the hook.</span>

<span class="sd">    .. note::</span>
<span class="sd">        See :ref:`backward-hooks-execution` for more information on how when this hook</span>
<span class="sd">        is executed, and how its execution is ordered relative to other hooks.</span>

<span class="sd">    Example::</span>

<span class="sd">        &gt;&gt;&gt; import torch</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; a = torch.rand(2, 3, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.rand(2, 3, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; c = a * b</span>
<span class="sd">        &gt;&gt;&gt; d = a * b</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; def fn(grads):</span>
<span class="sd">        ...     print([g is not None for g in grads])</span>
<span class="sd">        ...</span>
<span class="sd">        &gt;&gt;&gt; torch.autograd.graph.register_multi_grad_hook((a, b, c, d), fn)</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; c.sum().backward(retain_graph=True)</span>
<span class="sd">        [True, True, True, False]</span>
<span class="sd">        &gt;&gt;&gt; c.sum().backward(inputs=(a,), retain_graph=True)</span>
<span class="sd">        [True, False, True, False]</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">supported_modes</span> <span class="o">=</span> <span class="p">(</span><span class="s2">&quot;all&quot;</span><span class="p">,</span> <span class="s2">&quot;any&quot;</span><span class="p">)</span>
    <span class="n">lock</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Lock</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">mode</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">supported_modes</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Expects mode to be one of </span><span class="si">{</span><span class="n">supported_modes</span><span class="si">}</span><span class="s2"> but got </span><span class="si">{</span><span class="n">mode</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;all&quot;</span><span class="p">:</span>
        <span class="n">count</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="n">nb_calls</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">buffer</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]]]</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="n">grad_fns</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">_get_grad_fn_or_grad_acc</span><span class="p">,</span> <span class="n">tensors</span><span class="p">))</span>
        <span class="n">len_tensors</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>

        <span class="k">def</span> <span class="nf">get_inner_hook</span><span class="p">(</span><span class="n">idx</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="kc">None</span><span class="p">]:</span>
            <span class="k">def</span> <span class="nf">inner_hook</span><span class="p">(</span><span class="n">grad</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
                <span class="k">nonlocal</span> <span class="n">count</span><span class="p">,</span> <span class="n">nb_calls</span><span class="p">,</span> <span class="n">buffer</span><span class="p">,</span> <span class="n">fn</span>
                <span class="nb">id</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_current_graph_task_id</span><span class="p">()</span>
                <span class="k">assert</span> <span class="p">(</span>
                    <span class="nb">id</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span>
                <span class="p">),</span> <span class="s2">&quot;expected this hook to be called inside a backward call&quot;</span>
                <span class="n">count</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span> <span class="o">=</span> <span class="n">count</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="nb">id</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
                <span class="n">buffer</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span> <span class="o">=</span> <span class="n">buffer</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="nb">id</span><span class="p">,</span> <span class="p">[</span><span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">len_tensors</span><span class="p">)</span>

                <span class="k">with</span> <span class="n">lock</span><span class="p">:</span>
                    <span class="n">curr_count</span><span class="p">,</span> <span class="n">count</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span> <span class="o">=</span> <span class="n">count</span><span class="p">[</span><span class="nb">id</span><span class="p">],</span> <span class="n">count</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span>

                    <span class="k">if</span> <span class="n">curr_count</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                        <span class="c1"># On the first call, compute the actual nb_calls and buffer</span>
                        <span class="n">nb_calls</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span>
                            <span class="nb">map</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_will_engine_execute_node</span><span class="p">,</span> <span class="n">grad_fns</span><span class="p">)</span>
                        <span class="p">)</span>

                <span class="n">buffer</span><span class="p">[</span><span class="nb">id</span><span class="p">][</span><span class="n">idx</span><span class="p">]</span> <span class="o">=</span> <span class="n">grad</span>

                <span class="k">assert</span> <span class="n">nb_calls</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                <span class="k">if</span> <span class="n">curr_count</span> <span class="o">==</span> <span class="n">nb_calls</span> <span class="o">-</span> <span class="mi">1</span><span class="p">:</span>
                    <span class="n">fn</span> <span class="o">=</span> <span class="n">cast</span><span class="p">(</span><span class="n">Callable</span><span class="p">[[</span><span class="n">Sequence</span><span class="p">[</span><span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]]],</span> <span class="kc">None</span><span class="p">],</span> <span class="n">fn</span><span class="p">)</span>
                    <span class="n">fn</span><span class="p">(</span><span class="n">buffer</span><span class="p">[</span><span class="nb">id</span><span class="p">])</span>
                    <span class="k">del</span> <span class="n">count</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span>
                    <span class="k">del</span> <span class="n">buffer</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span>

            <span class="k">return</span> <span class="n">inner_hook</span>

        <span class="n">handles</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span>
            <span class="n">t</span><span class="o">.</span><span class="n">register_hook</span><span class="p">(</span><span class="n">get_inner_hook</span><span class="p">(</span><span class="n">i</span><span class="p">))</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>
        <span class="p">)</span>
    <span class="k">elif</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;any&quot;</span><span class="p">:</span>
        <span class="n">fn</span> <span class="o">=</span> <span class="n">cast</span><span class="p">(</span><span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">],</span> <span class="kc">None</span><span class="p">],</span> <span class="n">fn</span><span class="p">)</span>
        <span class="n">ran_hook</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">bool</span><span class="p">)</span>

        <span class="nd">@functools</span><span class="o">.</span><span class="n">wraps</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>
        <span class="k">def</span> <span class="nf">wrapped_fn</span><span class="p">(</span><span class="n">grad</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">nonlocal</span> <span class="n">ran_hook</span>
            <span class="nb">id</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_current_graph_task_id</span><span class="p">()</span>
            <span class="k">assert</span> <span class="nb">id</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="s2">&quot;expected this hook to be called inside a backward call&quot;</span>
            <span class="k">with</span> <span class="n">lock</span><span class="p">:</span>
                <span class="n">prev</span><span class="p">,</span> <span class="n">ran_hook</span><span class="p">[</span><span class="nb">id</span><span class="p">]</span> <span class="o">=</span> <span class="n">ran_hook</span><span class="p">[</span><span class="nb">id</span><span class="p">],</span> <span class="kc">True</span>
            <span class="k">if</span> <span class="n">prev</span><span class="p">:</span>
                <span class="k">return</span>
            <span class="n">fn</span><span class="p">(</span><span class="n">grad</span><span class="p">)</span>

        <span class="n">handles</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span>
            <span class="n">tensor</span><span class="o">.</span><span class="n">register_hook</span><span class="p">(</span><span class="n">wrapped_fn</span><span class="p">)</span>
            <span class="k">for</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="n">tensors</span>
            <span class="k">if</span> <span class="n">tensor</span><span class="o">.</span><span class="n">requires_grad</span>
        <span class="p">)</span>

    <span class="k">return</span> <span class="n">_MultiHandle</span><span class="p">(</span><span class="n">handles</span><span class="p">)</span>  <span class="c1"># type: ignore[possibly-undefined]</span></div>


<span class="c1"># NOTE [Allow mutation on tensors saved for backward]</span>
<span class="c1">#</span>
<span class="c1"># 1. Tensor gets saved for backward</span>
<span class="c1">#    - remember the python object id and the version of the tensor</span>
<span class="c1">#    - remember aliasing information (data_ptr of base + version)</span>
<span class="c1">#    - save the original so we control its lifetime</span>
<span class="c1"># 2. Any time a tensor gets in-placed</span>
<span class="c1">#    - for each tensor aliased to it:</span>
<span class="c1">#      - check using its object id and version to see if it has been saved</span>
<span class="c1">#      - if it has been saved, clone it</span>
<span class="c1">#      - delete the reference to the original</span>
<span class="c1"># 3. during backward</span>
<span class="c1">#    - if the clone exists, the tensor must&#39;ve been modified in-place</span>
<span class="n">_allow_mutation_on_saved_tensors_enabled</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>


<span class="n">_TID</span><span class="p">:</span> <span class="n">TypeAlias</span> <span class="o">=</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]</span>
<span class="n">_SID</span><span class="p">:</span> <span class="n">TypeAlias</span> <span class="o">=</span> <span class="n">Tuple</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">int</span><span class="p">]</span>


<span class="k">def</span> <span class="nf">_get_tid</span><span class="p">(</span><span class="n">tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">_TID</span><span class="p">:</span>
    <span class="c1"># FIXME: This is almost definitely a bug.</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span>
        <span class="n">tensor</span><span class="p">,</span>
        <span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_subclasses</span><span class="o">.</span><span class="n">fake_tensor</span><span class="o">.</span><span class="n">FakeTensor</span><span class="p">,</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_subclasses</span><span class="o">.</span><span class="n">functional_tensor</span><span class="o">.</span><span class="n">FunctionalTensor</span><span class="p">,</span>
        <span class="p">),</span>
    <span class="p">):</span>
        <span class="n">data_ptr</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">data_ptr</span> <span class="o">=</span> <span class="n">tensor</span><span class="o">.</span><span class="n">data_ptr</span><span class="p">()</span>
    <span class="k">return</span> <span class="p">(</span><span class="nb">id</span><span class="p">(</span><span class="n">tensor</span><span class="p">),</span> <span class="n">data_ptr</span><span class="p">,</span> <span class="n">tensor</span><span class="o">.</span><span class="n">_version</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_get_sid</span><span class="p">(</span><span class="n">tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">_SID</span><span class="p">:</span>
    <span class="c1"># FIXME: This is almost definitely a bug.</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span>
        <span class="n">tensor</span><span class="p">,</span>
        <span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_subclasses</span><span class="o">.</span><span class="n">fake_tensor</span><span class="o">.</span><span class="n">FakeTensor</span><span class="p">,</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_subclasses</span><span class="o">.</span><span class="n">functional_tensor</span><span class="o">.</span><span class="n">FunctionalTensor</span><span class="p">,</span>
        <span class="p">),</span>
    <span class="p">):</span>
        <span class="n">data_ptr</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">data_ptr</span> <span class="o">=</span> <span class="n">tensor</span><span class="o">.</span><span class="n">data_ptr</span><span class="p">()</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">data_ptr</span><span class="p">,</span> <span class="n">tensor</span><span class="o">.</span><span class="n">_version</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_Handle</span><span class="p">:</span>
    <span class="k">pass</span>


<span class="k">class</span> <span class="nc">_swap_with_cloned</span><span class="p">(</span><span class="n">saved_tensors_hooks</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ctx</span><span class="p">:</span> <span class="s2">&quot;_AllowMutationOnSavedContext&quot;</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">def</span> <span class="nf">pack_hook</span><span class="p">(</span><span class="n">tensor</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">_Handle</span><span class="p">:</span>
            <span class="n">tid</span> <span class="o">=</span> <span class="n">_get_tid</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
            <span class="n">sid</span> <span class="o">=</span> <span class="n">_get_sid</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span>
            <span class="c1"># Tensors saved for backward have an entry in _tid_to_weakhandle</span>
            <span class="n">handle</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">_Handle</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

            <span class="c1"># Save aliasing information</span>
            <span class="n">ctx</span><span class="o">.</span><span class="n">sid_to_tid</span><span class="p">[</span><span class="n">sid</span><span class="p">]</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">tid</span><span class="p">)</span>

            <span class="c1"># NB: The same tensor (of the same version) can be saved multiple times</span>
            <span class="k">if</span> <span class="n">tid</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="p">:</span>
                <span class="n">handle</span> <span class="o">=</span> <span class="n">_Handle</span><span class="p">()</span>
                <span class="n">ctx</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="p">[</span><span class="n">tid</span><span class="p">]</span> <span class="o">=</span> <span class="n">handle</span>
                <span class="n">ctx</span><span class="o">.</span><span class="n">original</span><span class="p">[</span><span class="n">handle</span><span class="p">]</span> <span class="o">=</span> <span class="n">tensor</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="c1"># Store an additional strong reference to the handle</span>
                <span class="n">handle</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="p">[</span><span class="n">tid</span><span class="p">]</span>
            <span class="k">return</span> <span class="n">handle</span>

        <span class="k">def</span> <span class="nf">unpack_hook</span><span class="p">(</span><span class="n">handle</span><span class="p">:</span> <span class="n">_Handle</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">:</span>
            <span class="n">error_msg</span> <span class="o">=</span> <span class="p">(</span>
                <span class="s2">&quot;Trying to backward outside of the &#39;allow_mutation_on_saved_tensors&#39; context&quot;</span>
                <span class="s2">&quot;in which the graph was originally recorded.&quot;</span>
            <span class="p">)</span>
            <span class="k">assert</span> <span class="n">_allow_mutation_on_saved_tensors_enabled</span><span class="p">,</span> <span class="n">error_msg</span>
            <span class="k">if</span> <span class="n">handle</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">cloned</span><span class="p">:</span>
                <span class="n">res</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">cloned</span><span class="p">[</span><span class="n">handle</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">assert</span> <span class="n">handle</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">original</span><span class="p">,</span> <span class="n">error_msg</span>
                <span class="n">res</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">original</span><span class="p">[</span><span class="n">handle</span><span class="p">]</span>
            <span class="k">return</span> <span class="n">res</span>

        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">(</span><span class="n">pack_hook</span><span class="p">,</span> <span class="n">unpack_hook</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_CloneArgBeforeMutateMode</span><span class="p">(</span><span class="n">TorchDispatchMode</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ctx</span><span class="p">:</span> <span class="s2">&quot;_AllowMutationOnSavedContext&quot;</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">ctx</span> <span class="o">=</span> <span class="n">ctx</span>

    <span class="k">def</span> <span class="nf">__torch_dispatch__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">func</span><span class="p">:</span> <span class="s2">&quot;OpOverload&quot;</span><span class="p">,</span>
        <span class="n">types</span><span class="p">:</span> <span class="n">Iterable</span><span class="p">[</span><span class="nb">type</span><span class="p">],</span>
        <span class="n">args</span><span class="p">:</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span> <span class="o">=</span> <span class="p">(),</span>
        <span class="n">kwargs</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
        <span class="n">kwargs</span> <span class="o">=</span> <span class="n">kwargs</span> <span class="ow">or</span> <span class="p">{}</span>

        <span class="k">for</span> <span class="n">idx</span><span class="p">,</span> <span class="n">arg</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">func</span><span class="o">.</span><span class="n">_schema</span><span class="o">.</span><span class="n">arguments</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">arg</span><span class="o">.</span><span class="n">alias_info</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">arg</span><span class="o">.</span><span class="n">alias_info</span><span class="o">.</span><span class="n">is_write</span><span class="p">:</span>
                <span class="n">t</span> <span class="o">=</span> <span class="n">kwargs</span><span class="p">[</span><span class="s2">&quot;out&quot;</span><span class="p">]</span> <span class="k">if</span> <span class="n">arg</span><span class="o">.</span><span class="n">is_out</span> <span class="k">else</span> <span class="n">args</span><span class="p">[</span><span class="n">idx</span><span class="p">]</span>
                <span class="n">tid</span> <span class="o">=</span> <span class="n">_get_tid</span><span class="p">(</span><span class="n">t</span><span class="p">)</span>
                <span class="n">sid</span> <span class="o">=</span> <span class="n">_get_sid</span><span class="p">(</span><span class="n">t</span><span class="p">)</span>
                <span class="n">ctx</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">ctx</span>
                <span class="k">if</span> <span class="n">sid</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">sid_to_tid</span><span class="p">:</span>
                    <span class="k">for</span> <span class="n">tid</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">sid_to_tid</span><span class="p">[</span><span class="n">sid</span><span class="p">]:</span>
                        <span class="k">if</span> <span class="n">tid</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="p">:</span>
                            <span class="c1"># We know that if tid is in sid_to_tid, then it must also be in</span>
                            <span class="c1"># tid_to_weakhandle. However, it is possible for the tensor to be</span>
                            <span class="c1"># saved at one point, but cleared by backward before it is modified</span>
                            <span class="c1"># in-place. Consider the following example:</span>
                            <span class="c1">#</span>
                            <span class="c1"># &gt;&gt;&gt; a = torch.randn(2, 3, requires_grad=True).clone()</span>
                            <span class="c1"># &gt;&gt;&gt; out = (a**2).sum()</span>
                            <span class="c1"># &gt;&gt;&gt; out.backward()</span>
                            <span class="c1"># &gt;&gt;&gt; a.sin_()</span>
                            <span class="k">continue</span>
                        <span class="n">handle</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="p">[</span><span class="n">tid</span><span class="p">]</span>
                        <span class="k">if</span> <span class="n">handle</span> <span class="ow">in</span> <span class="n">ctx</span><span class="o">.</span><span class="n">cloned</span><span class="p">:</span>
                            <span class="c1"># The same exact tensor has been cloned already</span>
                            <span class="k">continue</span>
                        <span class="n">ctx</span><span class="o">.</span><span class="n">cloned</span><span class="p">[</span><span class="n">handle</span><span class="p">]</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">original</span><span class="p">[</span><span class="n">handle</span><span class="p">]</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
                        <span class="k">del</span> <span class="n">ctx</span><span class="o">.</span><span class="n">original</span><span class="p">[</span><span class="n">handle</span><span class="p">]</span>

        <span class="k">return</span> <span class="n">func</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">_AllowMutationOnSavedContext</span><span class="p">:</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cloned</span><span class="p">:</span> <span class="n">MutableMapping</span><span class="p">[</span><span class="n">_Handle</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="n">WeakKeyDictionary</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">original</span><span class="p">:</span> <span class="n">MutableMapping</span><span class="p">[</span><span class="n">_Handle</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="n">WeakKeyDictionary</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="p">:</span> <span class="n">MutableMapping</span><span class="p">[</span><span class="n">_TID</span><span class="p">,</span> <span class="n">_Handle</span><span class="p">]</span> <span class="o">=</span> <span class="n">WeakValueDictionary</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sid_to_tid</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="n">_SID</span><span class="p">,</span> <span class="n">Set</span><span class="p">[</span><span class="n">_TID</span><span class="p">]]</span> <span class="o">=</span> <span class="n">defaultdict</span><span class="p">(</span><span class="nb">set</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">clear</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cloned</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">original</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">tid_to_weakhandle</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sid_to_tid</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>


<div class="viewcode-block" id="allow_mutation_on_saved_tensors"><a class="viewcode-back" href="../../../autograd.html#torch.autograd.graph.allow_mutation_on_saved_tensors">[docs]</a><span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
<span class="k">def</span> <span class="nf">allow_mutation_on_saved_tensors</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="p">(</span>
    <span class="n">Generator</span><span class="p">[</span><span class="n">_AllowMutationOnSavedContext</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Context manager under which mutating tensors saved for backward is allowed.</span>

<span class="sd">    Under this context manager, tensors saved for backward are cloned on mutation,</span>
<span class="sd">    so the original version can still be used during backward. Normally, mutating a tensor</span>
<span class="sd">    saved for backward will result in an error raised when it&#39;s used during backward.</span>

<span class="sd">    To ensure the correct behavior, both the forward and backward should be run under</span>
<span class="sd">    the same context manager.</span>

<span class="sd">    Returns:</span>
<span class="sd">        An _AllowMutationOnSavedContext object storing the state managed by this</span>
<span class="sd">        context manager. This object can be useful for debugging purposes. The state</span>
<span class="sd">        managed by the context manager is automatically cleared upon exiting.</span>

<span class="sd">    Example::</span>

<span class="sd">        &gt;&gt;&gt; import torch</span>
<span class="sd">        &gt;&gt;&gt; with torch.autograd.graph.allow_mutation_on_saved_tensors():</span>
<span class="sd">        ...     # forward</span>
<span class="sd">        ...     a = torch.ones(2, 3, requires_grad=True)</span>
<span class="sd">        ...     b = a.clone()</span>
<span class="sd">        ...     out = (b**2).sum()</span>
<span class="sd">        ...     b.sin_()</span>
<span class="sd">        ...     # backward</span>
<span class="sd">        ...     out.sum().backward()</span>
<span class="sd">        ...</span>
<span class="sd">        tensor([[0.8415, 0.8415, 0.8415],</span>
<span class="sd">                [0.8415, 0.8415, 0.8415]], grad_fn=&lt;SinBackward0&gt;)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_allow_mutation_on_saved_tensors_enabled</span>

    <span class="n">ctx</span> <span class="o">=</span> <span class="n">_AllowMutationOnSavedContext</span><span class="p">()</span>

    <span class="k">with</span> <span class="n">_swap_with_cloned</span><span class="p">(</span><span class="n">ctx</span><span class="p">),</span> <span class="n">_CloneArgBeforeMutateMode</span><span class="p">(</span><span class="n">ctx</span><span class="p">):</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">_allow_mutation_on_saved_tensors_enabled</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                    <span class="s2">&quot;allow_mutation_on_saved_tensors contexts cannot be nested&quot;</span>
                <span class="p">)</span>
            <span class="n">_allow_mutation_on_saved_tensors_enabled</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="k">yield</span> <span class="n">ctx</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="n">ctx</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
            <span class="n">_allow_mutation_on_saved_tensors_enabled</span> <span class="o">=</span> <span class="kc">False</span></div>


<span class="k">def</span> <span class="nf">_register_logging_hooks_on_whole_graph</span><span class="p">(</span>
    <span class="n">t_outputs</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">GradientEdge</span><span class="p">]],</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Callable</span><span class="p">[[],</span> <span class="kc">None</span><span class="p">]:</span>
    <span class="n">grad_fns</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="n">_get_grad_fn_or_grad_acc</span><span class="p">,</span> <span class="n">t_outputs</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">iter_graph</span><span class="p">(</span><span class="n">roots</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Node</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Iterator</span><span class="p">[</span><span class="n">Node</span><span class="p">]:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">roots</span><span class="p">:</span>
            <span class="k">return</span>
        <span class="n">seen</span><span class="p">:</span> <span class="n">Set</span><span class="p">[</span><span class="n">Node</span><span class="p">]</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
        <span class="n">q</span><span class="p">:</span> <span class="n">Deque</span><span class="p">[</span><span class="n">Node</span><span class="p">]</span> <span class="o">=</span> <span class="n">deque</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">roots</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">node</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">seen</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>
                <span class="n">q</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>

        <span class="k">while</span> <span class="n">q</span><span class="p">:</span>
            <span class="n">node</span> <span class="o">=</span> <span class="n">q</span><span class="o">.</span><span class="n">popleft</span><span class="p">()</span>
            <span class="k">for</span> <span class="n">fn</span><span class="p">,</span> <span class="n">_</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">next_functions</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">fn</span> <span class="ow">in</span> <span class="n">seen</span> <span class="ow">or</span> <span class="n">fn</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="k">continue</span>
                <span class="n">seen</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>
                <span class="n">q</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>

            <span class="k">yield</span> <span class="n">node</span>

    <span class="k">def</span> <span class="nf">fmt</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
        <span class="c1"># Avoid circular import</span>
        <span class="kn">from</span> <span class="nn">torch.testing._internal.common_utils</span> <span class="kn">import</span> <span class="n">dtype_abbrs</span>

        <span class="k">if</span> <span class="n">t</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span> <span class="s2">&quot;None&quot;</span>
        <span class="k">return</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">dtype_abbrs</span><span class="p">[</span><span class="n">t</span><span class="o">.</span><span class="n">dtype</span><span class="p">]</span><span class="si">}</span><span class="s2">[</span><span class="si">{</span><span class="s1">&#39;, &#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="nb">map</span><span class="p">(</span><span class="nb">str</span><span class="p">,</span><span class="w"> </span><span class="n">t</span><span class="o">.</span><span class="n">shape</span><span class="p">))</span><span class="si">}</span><span class="s2">]&quot;</span>

    <span class="k">def</span> <span class="nf">prehook</span><span class="p">(</span><span class="n">grad_outputs</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">node</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_current_autograd_node</span><span class="p">()</span>
        <span class="n">grad_outputs_str</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;[</span><span class="si">{</span><span class="s1">&#39;,&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">fmt</span><span class="p">(</span><span class="n">t</span><span class="p">)</span><span class="w"> </span><span class="k">for</span><span class="w"> </span><span class="n">t</span><span class="w"> </span><span class="ow">in</span><span class="w"> </span><span class="n">grad_outputs</span><span class="p">)</span><span class="si">}</span><span class="s2">]&quot;</span>
        <span class="n">log_str</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;Executing: </span><span class="si">{</span><span class="n">node</span><span class="si">}</span><span class="s2"> with grad_outputs: </span><span class="si">{</span><span class="n">grad_outputs_str</span><span class="si">}</span><span class="s2">&quot;</span>
        <span class="n">log</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="n">log_str</span><span class="p">)</span>

    <span class="n">handles</span> <span class="o">=</span> <span class="p">[</span><span class="n">node</span><span class="o">.</span><span class="n">register_prehook</span><span class="p">(</span><span class="n">prehook</span><span class="p">)</span> <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">iter_graph</span><span class="p">(</span><span class="n">grad_fns</span><span class="p">)]</span>

    <span class="k">def</span> <span class="nf">unregister_hooks</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">handle</span> <span class="ow">in</span> <span class="n">handles</span><span class="p">:</span>
            <span class="n">handle</span><span class="o">.</span><span class="n">remove</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">unregister_hooks</span>


<span class="k">def</span> <span class="nf">_engine_run_backward</span><span class="p">(</span>
    <span class="n">t_outputs</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">GradientEdge</span><span class="p">]],</span>
    <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span>
    <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">...</span><span class="p">]:</span>
    <span class="n">attach_logging_hooks</span> <span class="o">=</span> <span class="n">log</span><span class="o">.</span><span class="n">getEffectiveLevel</span><span class="p">()</span> <span class="o">&lt;=</span> <span class="n">logging</span><span class="o">.</span><span class="n">DEBUG</span>
    <span class="k">if</span> <span class="n">attach_logging_hooks</span><span class="p">:</span>
        <span class="n">unregister_hooks</span> <span class="o">=</span> <span class="n">_register_logging_hooks_on_whole_graph</span><span class="p">(</span><span class="n">t_outputs</span><span class="p">)</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">Variable</span><span class="o">.</span><span class="n">_execution_engine</span><span class="o">.</span><span class="n">run_backward</span><span class="p">(</span>  <span class="c1"># Calls into the C++ engine to run the backward pass</span>
            <span class="n">t_outputs</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span>
        <span class="p">)</span>  <span class="c1"># Calls into the C++ engine to run the backward pass</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">attach_logging_hooks</span><span class="p">:</span>
            <span class="n">unregister_hooks</span><span class="p">()</span>  <span class="c1"># type: ignore[possibly-undefined]</span>
</pre></div>

             </article>
             
            </div>
            <footer>
  

    <hr>

  
  <div role="contentinfo">
    <p>
        &copy; Copyright 2024, PyTorch Contributors.

    </p>
  </div>
    
      <div>
        Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
      </div>
     

</footer>

          </div>
<script>

var match = window.location.href.match(/\/_[a-zA-Z0-9_]*.html|_dynamo/gi);
var url = window.location.href.lastIndexOf(match[match.length-1]);

if (url)
  {
    var div = '<div class="admonition note"><p class="admonition-title">Note</p><p><i class="fa fa-exclamation-circle" aria-hidden="true">&nbsp</i> This page describes an internal API which is not intended to be used outside of the PyTorch codebase and can be modified or removed without notice.</p></div>'
    document.getElementById("pytorch-article").insertAdjacentHTML('afterBegin', div)
  }
</script>
        </div>

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              
            </div>
          </div>
        </div>
      </section>
    </div>

  
       <script type="text/javascript" id="documentation_options" data-url_root="../../../" src="../../../_static/documentation_options.js"></script>
         <script data-url_root="../../../" id="documentation_options" src="../../../_static/documentation_options.js"></script>
         <script src="../../../_static/jquery.js"></script>
         <script src="../../../_static/underscore.js"></script>
         <script src="../../../_static/_sphinx_javascript_frameworks_compat.js"></script>
         <script src="../../../_static/doctools.js"></script>
         <script src="../../../_static/sphinx_highlight.js"></script>
         <script src="../../../_static/clipboard.min.js"></script>
         <script src="../../../_static/copybutton.js"></script>
     

  <script type="text/javascript" src="../../../_static/js/vendor/popper.min.js"></script>
  <script type="text/javascript" src="../../../_static/js/vendor/bootstrap.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/list.js/1.5.0/list.min.js"></script>
  <script type="text/javascript" src="../../../_static/js/theme.js"></script>

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>
 
<script script type="text/javascript">
  var collapsedSections = ['Developer Notes', 'Language Bindings', 'Libraries', 'Community'];
</script>

<img height="1" width="1" style="border-style:none;" alt="" src="https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&amp;guid=ON&amp;script=0"/>


  <!-- Begin Footer -->

  <div class="container-fluid docs-tutorials-resources" id="docs-tutorials-resources">
    <div class="container">
      <div class="row">
        <div class="col-md-4 text-center">
          <h2>Docs</h2>
          <p>Access comprehensive developer documentation for PyTorch</p>
          <a class="with-right-arrow" href="https://pytorch.org/docs/stable/index.html">View Docs</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Tutorials</h2>
          <p>Get in-depth tutorials for beginners and advanced developers</p>
          <a class="with-right-arrow" href="https://pytorch.org/tutorials">View Tutorials</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Resources</h2>
          <p>Find development resources and get your questions answered</p>
          <a class="with-right-arrow" href="https://pytorch.org/resources">View Resources</a>
        </div>
      </div>
    </div>
  </div>

  <footer class="site-footer">
    <div class="container footer-container">
      <div class="footer-logo-wrapper">
        <a href="https://pytorch.org/" class="footer-logo"></a>
      </div>

      <div class="footer-links-wrapper">
        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/">PyTorch</a></li>
            <li><a href="https://pytorch.org/get-started">Get Started</a></li>
            <li><a href="https://pytorch.org/features">Features</a></li>
            <li><a href="https://pytorch.org/ecosystem">Ecosystem</a></li>
            <li><a href="https://pytorch.org/blog/">Blog</a></li>
            <li><a href="https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md">Contributing</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/resources">Resources</a></li>
            <li><a href="https://pytorch.org/tutorials">Tutorials</a></li>
            <li><a href="https://pytorch.org/docs/stable/index.html">Docs</a></li>
            <li><a href="https://discuss.pytorch.org" target="_blank">Discuss</a></li>
            <li><a href="https://github.com/pytorch/pytorch/issues" target="_blank">Github Issues</a></li>
            <li><a href="https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf" target="_blank">Brand Guidelines</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title">Stay up to date</li>
            <li><a href="https://www.facebook.com/pytorch" target="_blank">Facebook</a></li>
            <li><a href="https://twitter.com/pytorch" target="_blank">Twitter</a></li>
            <li><a href="https://www.youtube.com/pytorch" target="_blank">YouTube</a></li>
            <li><a href="https://www.linkedin.com/company/pytorch" target="_blank">LinkedIn</a></li>
          </ul>  
          </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title">PyTorch Podcasts</li>
            <li><a href="https://open.spotify.com/show/6UzHKeiy368jKfQMKKvJY5" target="_blank">Spotify</a></li>
            <li><a href="https://podcasts.apple.com/us/podcast/pytorch-developer-podcast/id1566080008" target="_blank">Apple</a></li>
            <li><a href="https://www.google.com/podcasts?feed=aHR0cHM6Ly9mZWVkcy5zaW1wbGVjYXN0LmNvbS9PQjVGa0lsOA%3D%3D" target="_blank">Google</a></li>
            <li><a href="https://music.amazon.com/podcasts/7a4e6f0e-26c2-49e9-a478-41bd244197d0/PyTorch-Developer-Podcast?" target="_blank">Amazon</a></li>
          </ul>
         </div>
        </div>
        
        <div class="privacy-policy">
          <ul>
            <li class="privacy-policy-links"><a href="https://www.linuxfoundation.org/terms/" target="_blank">Terms</a></li>
            <li class="privacy-policy-links">|</li>
            <li class="privacy-policy-links"><a href="https://www.linuxfoundation.org/privacy-policy/" target="_blank">Privacy</a></li>
          </ul>
        </div>
        <div class="copyright">
        <p>© Copyright The Linux Foundation. The PyTorch Foundation is a project of The Linux Foundation.
          For web site terms of use, trademark policy and other policies applicable to The PyTorch Foundation please see
          <a href="https://www.linuxfoundation.org/policies/">www.linuxfoundation.org/policies/</a>. The PyTorch Foundation supports the PyTorch open source
          project, which has been established as PyTorch Project a Series of LF Projects, LLC. For policies applicable to the PyTorch Project a Series of LF Projects, LLC,
          please see <a href="https://www.lfprojects.org/policies/">www.lfprojects.org/policies/</a>.</p>
      </div>
     </div>

  </footer>

  <div class="cookie-banner-wrapper">
  <div class="container">
    <p class="gdpr-notice">To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: <a href="https://www.facebook.com/policies/cookies/">Cookies Policy</a>.</p>
    <img class="close-button" src="../../../_static/images/pytorch-x.svg">
  </div>
</div>

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

  <div class="mobile-main-menu">
    <div class="container-fluid">
      <div class="container">
        <div class="mobile-main-menu-header-container">
          <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
          <a class="main-menu-close-button" href="#" data-behavior="close-mobile-menu"></a>
        </div>
      </div>
    </div>

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
           <li class="resources-mobile-menu-title">
             <a>Learn</a>
           </li>
           <ul class="resources-mobile-menu-items">
             <li>
               <a href="https://pytorch.org/get-started">Get Started</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials">Tutorials</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials/beginner/basics/intro.html">Learn the Basics</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials/recipes/recipes_index.html">PyTorch Recipes</a>
             </li>
             <li>
               <a href="https://pytorch.org/tutorials/beginner/introyt.html">Introduction to PyTorch - YouTube Series</a>
             </li>
           </ul>
           <li class="resources-mobile-menu-title">
             <a>Ecosystem</a>
           </li>
           <ul class="resources-mobile-menu-items">
             <li>
               <a href="https://pytorch.org/ecosystem">Tools</a>
             </li>
             <li>
               <a href="https://pytorch.org/#community-module">Community</a>
             </li>
             <li>
               <a href="https://discuss.pytorch.org/">Forums</a>
             </li>
             <li>
               <a href="https://pytorch.org/resources">Developer Resources</a>
             </li>
             <li>
               <a href="https://pytorch.org/ecosystem/contributor-awards-2023">Contributor Awards - 2023</a>
             </li>
           </ul>

           <li class="resources-mobile-menu-title">
             <a>Edge</a>
           </li>

           <ul class="resources-mobile-menu-items">
             <li>
               <a href="https://pytorch.org/edge">About PyTorch Edge</a>
             </li>
             
             <li>
               <a href="https://pytorch.org/executorch-overview">ExecuTorch</a>
             </li>
           </ul>

           <li class="resources-mobile-menu-title">
             <a>Docs</a>
           </li>

           <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/docs/stable/index.html">PyTorch</a>
            </li>

            <li>
              <a href="https://pytorch.org/pytorch-domains">PyTorch Domains</a>
            </li>
          </ul>

          <li class="resources-mobile-menu-title">
            <a>Blog & News</a>
          </li>
            
           <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/blog/">PyTorch Blog</a>
            </li>
            <li>
              <a href="https://pytorch.org/community-blog">Community Blog</a>
            </li>

            <li>
              <a href="https://pytorch.org/videos">Videos</a>
            </li>

            <li>
              <a href="https://pytorch.org/community-stories">Community Stories</a>
            </li>
            <li>
              <a href="https://pytorch.org/events">Events</a>
            </li>
          </ul>
          
          <li class="resources-mobile-menu-title">
            <a>About</a>
          </li>

          <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/foundation">PyTorch Foundation</a>
            </li>
            <li>
              <a href="https://pytorch.org/governing-board">Governing Board</a>
            </li>
          </ul>
        </ul>
      </div>
    </div>
  </div>

  <!-- End Mobile Menu -->

  <script type="text/javascript" src="../../../_static/js/vendor/anchor.min.js"></script>

  <script type="text/javascript">
    $(document).ready(function() {
      mobileMenu.bind();
      mobileTOC.bind();
      pytorchAnchors.bind();
      sideMenus.bind();
      scrollToAnchor.bind();
      highlightNavigation.bind();
      mainMenuDropdown.bind();
      filterTags.bind();

      // Add class to links that have code blocks, since we cannot create links in code blocks
      $("article.pytorch-article a span.pre").each(function(e) {
        $(this).closest("a").addClass("has-code");
      });
    })
  </script>
</body>
</html>