2.6/export.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.export &mdash; PyTorch 2.6 documentation</title>
  

    <link rel="canonical" href="https://pytorch.org/docs/stable/export.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" />
    <link rel="next" title="torch.export IR Specification" href="export.ir_spec.html" />
    <link rel="prev" title="torch.backends" href="backends.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 class="current">
<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 current"><a class="current reference internal" href="#">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>torch.export</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
            
            <a href="_sources/export.rst.txt" rel="nofollow"><img src="_static/images/view-page-source-icon.svg"></a>
          
        
      </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">
              
  <div class="section" id="torch-export">
<span id="id1"></span><h1>torch.export<a class="headerlink" href="#torch-export" title="Permalink to this heading">¶</a></h1>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This feature is a prototype under active development and there WILL BE
BREAKING CHANGES in the future.</p>
</div>
<div class="section" id="overview">
<h2>Overview<a class="headerlink" href="#overview" title="Permalink to this heading">¶</a></h2>
<p><a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.export()</span></code></a> takes an arbitrary Python callable (a
<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.Module</span></code></a>, a function or a method) and produces a traced graph
representing only the Tensor computation of the function in an Ahead-of-Time
(AOT) fashion, which can subsequently be executed with different outputs or
serialized.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.export</span> <span class="kn">import</span> <span class="n">export</span>

<span class="k">class</span> <span class="nc">Mod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</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">y</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">a</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sin</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">b</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cos</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">a</span> <span class="o">+</span> <span class="n">b</span>

<span class="n">example_args</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">))</span>

<span class="n">exported_program</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">ExportedProgram</span> <span class="o">=</span> <span class="n">export</span><span class="p">(</span>
    <span class="n">Mod</span><span class="p">(),</span> <span class="n">args</span><span class="o">=</span><span class="n">example_args</span>
<span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">exported_program</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
    <span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
        <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[10, 10]&quot;</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="s2">&quot;f32[10, 10]&quot;</span><span class="p">):</span>
            <span class="c1"># code: a = torch.sin(x)</span>
            <span class="n">sin</span><span class="p">:</span> <span class="s2">&quot;f32[10, 10]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">sin</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>

            <span class="c1"># code: b = torch.cos(y)</span>
            <span class="n">cos</span><span class="p">:</span> <span class="s2">&quot;f32[10, 10]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">cos</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>

            <span class="c1"># code: return a + b</span>
            <span class="n">add</span><span class="p">:</span> <span class="n">f32</span><span class="p">[</span><span class="mi">10</span><span class="p">,</span> <span class="mi">10</span><span class="p">]</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">sin</span><span class="p">,</span> <span class="n">cos</span><span class="p">)</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">add</span><span class="p">,)</span>

    <span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
        <span class="n">ExportGraphSignature</span><span class="p">(</span>
            <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
                <span class="n">InputSpec</span><span class="p">(</span>
                    <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                    <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x&#39;</span><span class="p">),</span>
                    <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                    <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
                <span class="p">),</span>
                <span class="n">InputSpec</span><span class="p">(</span>
                    <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                    <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;y&#39;</span><span class="p">),</span>
                    <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                    <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
                <span class="p">)</span>
            <span class="p">],</span>
            <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
                <span class="n">OutputSpec</span><span class="p">(</span>
                    <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                    <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add&#39;</span><span class="p">),</span>
                    <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
                <span class="p">)</span>
            <span class="p">]</span>
        <span class="p">)</span>
    <span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{}</span>
</pre></div>
</div>
<p><code class="docutils literal notranslate"><span class="pre">torch.export</span></code> produces a clean intermediate representation (IR) with the
following invariants. More specifications about the IR can be found
<a class="reference internal" href="export.ir_spec.html#export-ir-spec"><span class="std std-ref">here</span></a>.</p>
<ul class="simple">
<li><p><strong>Soundness</strong>: It is guaranteed to be a sound representation of the original
program, and maintains the same calling conventions of the original program.</p></li>
<li><p><strong>Normalized</strong>: There are no Python semantics within the graph. Submodules
from the original programs are inlined to form one fully flattened
computational graph.</p></li>
<li><p><strong>Graph properties</strong>: The graph is purely functional, meaning it does not
contain operations with side effects such as mutations or aliasing. It does
not mutate any intermediate values, parameters, or buffers.</p></li>
<li><p><strong>Metadata</strong>: The graph contains metadata captured during tracing, such as a
stacktrace from user’s code.</p></li>
</ul>
<p>Under the hood, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> leverages the following latest technologies:</p>
<ul class="simple">
<li><p><strong>TorchDynamo (torch._dynamo)</strong> is an internal API that uses a CPython feature
called the Frame Evaluation API to safely trace PyTorch graphs. This
provides a massively improved graph capturing experience, with much fewer
rewrites needed in order to fully trace the PyTorch code.</p></li>
<li><p><strong>AOT Autograd</strong> provides a functionalized PyTorch graph and ensures the graph
is decomposed/lowered to the ATen operator set.</p></li>
<li><p><strong>Torch FX (torch.fx)</strong> is the underlying representation of the graph,
allowing flexible Python-based transformations.</p></li>
</ul>
<div class="section" id="existing-frameworks">
<h3>Existing frameworks<a class="headerlink" href="#existing-frameworks" title="Permalink to this heading">¶</a></h3>
<p><a class="reference internal" href="generated/torch.compile.html#torch.compile" title="torch.compile"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.compile()</span></code></a> also utilizes the same PT2 stack as <code class="docutils literal notranslate"><span class="pre">torch.export</span></code>, but
is slightly different:</p>
<ul class="simple">
<li><p><strong>JIT vs. AOT</strong>: <a class="reference internal" href="generated/torch.compile.html#torch.compile" title="torch.compile"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.compile()</span></code></a> is a JIT compiler whereas
which is not intended to be used to produce compiled artifacts outside of
deployment.</p></li>
<li><p><strong>Partial vs. Full Graph Capture</strong>: When <a class="reference internal" href="generated/torch.compile.html#torch.compile" title="torch.compile"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.compile()</span></code></a> runs into an
untraceable part of a model, it will “graph break” and fall back to running
the program in the eager Python runtime. In comparison, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> aims
to get a full graph representation of a PyTorch model, so it will error out
when something untraceable is reached. Since <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> produces a full
graph disjoint from any Python features or runtime, this graph can then be
saved, loaded, and run in different environments and languages.</p></li>
<li><p><strong>Usability tradeoff</strong>: Since <a class="reference internal" href="generated/torch.compile.html#torch.compile" title="torch.compile"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.compile()</span></code></a> is able to fallback to the
Python runtime whenever it reaches something untraceable, it is a lot more
flexible. <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> will instead require users to provide more
information or rewrite their code to make it traceable.</p></li>
</ul>
<p>Compared to <a class="reference internal" href="fx.html#torch.fx.symbolic_trace" title="torch.fx.symbolic_trace"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.fx.symbolic_trace()</span></code></a>, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> traces using
TorchDynamo which operates at the Python bytecode level, giving it the ability
to trace arbitrary Python constructs not limited by what Python operator
overloading supports. Additionally, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> keeps fine-grained track of
tensor metadata, so that conditionals on things like tensor shapes do not
fail tracing. In general, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> is expected to work on more user
programs, and produce lower-level graphs (at the <code class="docutils literal notranslate"><span class="pre">torch.ops.aten</span></code> operator
level). Note that users can still use <a class="reference internal" href="fx.html#torch.fx.symbolic_trace" title="torch.fx.symbolic_trace"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.fx.symbolic_trace()</span></code></a> as a
preprocessing step before <code class="docutils literal notranslate"><span class="pre">torch.export</span></code>.</p>
<p>Compared to <a class="reference internal" href="generated/torch.jit.script.html#torch.jit.script" title="torch.jit.script"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.jit.script()</span></code></a>, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> does not capture Python
control flow or data structures, but it supports more Python language features
than TorchScript (as it is easier to have comprehensive coverage over Python
bytecodes). The resulting graphs are simpler and only have straight line control
flow (except for explicit control flow operators).</p>
<p>Compared to <a class="reference internal" href="generated/torch.jit.trace.html#torch.jit.trace" title="torch.jit.trace"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.jit.trace()</span></code></a>, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> is sound: it is able to
trace code that performs integer computation on sizes and records all of the
side-conditions necessary to show that a particular trace is valid for other
inputs.</p>
</div>
</div>
<div class="section" id="exporting-a-pytorch-model">
<h2>Exporting a PyTorch Model<a class="headerlink" href="#exporting-a-pytorch-model" title="Permalink to this heading">¶</a></h2>
<div class="section" id="an-example">
<h3>An Example<a class="headerlink" href="#an-example" title="Permalink to this heading">¶</a></h3>
<p>The main entrypoint is through <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.export()</span></code></a>, which takes a
callable (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.Module</span></code></a>, function, or method) and sample inputs, and
captures the computation graph into an <a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.export.ExportedProgram</span></code></a>. An
example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.export</span> <span class="kn">import</span> <span class="n">export</span>

<span class="c1"># Simple module for demonstration</span>
<span class="k">class</span> <span class="nc">M</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</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="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">conv</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">(</span>
            <span class="n">in_channels</span><span class="o">=</span><span class="mi">3</span><span class="p">,</span> <span class="n">out_channels</span><span class="o">=</span><span class="mi">16</span><span class="p">,</span> <span class="n">kernel_size</span><span class="o">=</span><span class="mi">3</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="mi">1</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">relu</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ReLU</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">maxpool</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">MaxPool2d</span><span class="p">(</span><span class="n">kernel_size</span><span class="o">=</span><span class="mi">3</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</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">constant</span><span class="o">=</span><span class="kc">None</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">a</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">conv</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">a</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="n">constant</span><span class="p">)</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">maxpool</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="n">a</span><span class="p">))</span>

<span class="n">example_args</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">256</span><span class="p">,</span> <span class="mi">256</span><span class="p">),)</span>
<span class="n">example_kwargs</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;constant&quot;</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">16</span><span class="p">,</span> <span class="mi">256</span><span class="p">,</span> <span class="mi">256</span><span class="p">)}</span>

<span class="n">exported_program</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">ExportedProgram</span> <span class="o">=</span> <span class="n">export</span><span class="p">(</span>
    <span class="n">M</span><span class="p">(),</span> <span class="n">args</span><span class="o">=</span><span class="n">example_args</span><span class="p">,</span> <span class="n">kwargs</span><span class="o">=</span><span class="n">example_kwargs</span>
<span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">exported_program</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
    <span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">:</span> <span class="s2">&quot;f32[16, 3, 3, 3]&quot;</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">:</span> <span class="s2">&quot;f32[16]&quot;</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 256, 256]&quot;</span><span class="p">,</span> <span class="n">constant</span><span class="p">:</span> <span class="s2">&quot;f32[1, 16, 256, 256]&quot;</span><span class="p">):</span>
            <span class="c1"># code: a = self.conv(x)</span>
            <span class="n">conv2d</span><span class="p">:</span> <span class="s2">&quot;f32[1, 16, 256, 256]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">conv2d</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">,</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>

            <span class="c1"># code: a.add_(constant)</span>
            <span class="n">add_</span><span class="p">:</span> <span class="s2">&quot;f32[1, 16, 256, 256]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add_</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">conv2d</span><span class="p">,</span> <span class="n">constant</span><span class="p">)</span>

            <span class="c1"># code: return self.maxpool(self.relu(a))</span>
            <span class="n">relu</span><span class="p">:</span> <span class="s2">&quot;f32[1, 16, 256, 256]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">relu</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">add_</span><span class="p">)</span>
            <span class="n">max_pool2d</span><span class="p">:</span> <span class="s2">&quot;f32[1, 16, 85, 85]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">max_pool2d</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">relu</span><span class="p">,</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">],</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">])</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">max_pool2d</span><span class="p">,)</span>

<span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
    <span class="n">ExportGraphSignature</span><span class="p">(</span>
        <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;constant&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">],</span>
        <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;max_pool2d&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">]</span>
    <span class="p">)</span>
<span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{}</span>
</pre></div>
</div>
<p>Inspecting the <code class="docutils literal notranslate"><span class="pre">ExportedProgram</span></code>, we can note the following:</p>
<ul class="simple">
<li><p>The <a class="reference internal" href="fx.html#torch.fx.Graph" title="torch.fx.Graph"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.fx.Graph</span></code></a> contains the computation graph of the original
program, along with records of the original code for easy debugging.</p></li>
<li><p>The graph contains only <code class="docutils literal notranslate"><span class="pre">torch.ops.aten</span></code> operators found <a class="reference external" href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml">here</a>
and custom operators, and is fully functional, without any inplace operators
such as <code class="docutils literal notranslate"><span class="pre">torch.add_</span></code>.</p></li>
<li><p>The parameters (weight and bias to conv) are lifted as inputs to the graph,
resulting in no <code class="docutils literal notranslate"><span class="pre">get_attr</span></code> nodes in the graph, which previously existed in
the result of <a class="reference internal" href="fx.html#torch.fx.symbolic_trace" title="torch.fx.symbolic_trace"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.fx.symbolic_trace()</span></code></a>.</p></li>
<li><p>The <a class="reference internal" href="#torch.export.ExportGraphSignature" title="torch.export.ExportGraphSignature"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.export.ExportGraphSignature</span></code></a> models the input and output
signature, along with specifying which inputs are parameters.</p></li>
<li><p>The resulting shape and dtype of tensors produced by each node in the graph is
noted. For example, the <code class="docutils literal notranslate"><span class="pre">convolution</span></code> node will result in a tensor of dtype
<code class="docutils literal notranslate"><span class="pre">torch.float32</span></code> and shape (1, 16, 256, 256).</p></li>
</ul>
</div>
<div class="section" id="non-strict-export">
<span id="id2"></span><h3>Non-Strict Export<a class="headerlink" href="#non-strict-export" title="Permalink to this heading">¶</a></h3>
<p>In PyTorch 2.3, we introduced a new mode of tracing called <strong>non-strict mode</strong>.
It’s still going through hardening, so if you run into any issues, please file
them to Github with the “oncall: export” tag.</p>
<p>In <em>non-strict mode</em>, we trace through the program using the Python interpreter.
Your code will execute exactly as it would in eager mode; the only difference is
that all Tensor objects will be replaced by ProxyTensors, which will record all
their operations into a graph.</p>
<p>In <em>strict</em> mode, which is currently the default, we first trace through the
program using TorchDynamo, a bytecode analysis engine. TorchDynamo does not
actually execute your Python code. Instead, it symbolically analyzes it and
builds a graph based on the results. This analysis allows torch.export to
provide stronger guarantees about safety, but not all Python code is supported.</p>
<p>An example of a case where one might want to use non-strict mode is if you run
into a unsupported TorchDynamo feature that might not be easily solved, and you
know the python code is not exactly needed for computation. For example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">contextlib</span>
<span class="kn">import</span> <span class="nn">torch</span>

<span class="k">class</span> <span class="nc">ContextManager</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="bp">self</span><span class="o">.</span><span class="n">count</span> <span class="o">=</span> <span class="mi">0</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="bp">self</span><span class="o">.</span><span class="n">count</span> <span class="o">+=</span> <span class="mi">1</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="n">exc_type</span><span class="p">,</span> <span class="n">exc_value</span><span class="p">,</span> <span class="n">traceback</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">count</span> <span class="o">-=</span> <span class="mi">1</span>

<span class="k">class</span> <span class="nc">M</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">with</span> <span class="n">ContextManager</span><span class="p">():</span>
            <span class="k">return</span> <span class="n">x</span><span class="o">.</span><span class="n">sin</span><span class="p">()</span> <span class="o">+</span> <span class="n">x</span><span class="o">.</span><span class="n">cos</span><span class="p">()</span>

<span class="n">export</span><span class="p">(</span><span class="n">M</span><span class="p">(),</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">),),</span> <span class="n">strict</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>  <span class="c1"># Non-strict traces successfully</span>
<span class="n">export</span><span class="p">(</span><span class="n">M</span><span class="p">(),</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">),))</span>  <span class="c1"># Strict mode fails with torch._dynamo.exc.Unsupported: ContextManager</span>
</pre></div>
</div>
<p>In this example, the first call using non-strict mode (through the
<code class="docutils literal notranslate"><span class="pre">strict=False</span></code> flag) traces successfully whereas the second call using strict
mode (default) results with a failure, where TorchDynamo is unable to support
context managers. One option is to rewrite the code (see <a class="reference internal" href="#limitations-of-torch-export"><span class="std std-ref">Limitations of torch.export</span></a>), but seeing as the context manager does not affect the tensor
computations in the model, we can go with the non-strict mode’s result.</p>
</div>
<div class="section" id="export-for-training-and-inference">
<span id="training-export"></span><h3>Export for Training and Inference<a class="headerlink" href="#export-for-training-and-inference" title="Permalink to this heading">¶</a></h3>
<p>In PyTorch 2.5, we introduced a new API called <code class="xref py py-func docutils literal notranslate"><span class="pre">export_for_training()</span></code>.
It’s still going through hardening, so if you run into any issues, please file
them to Github with the “oncall: export” tag.</p>
<p>In this API, we produce the most generic IR that contains all ATen operators
(including both functional and non-functional) which can be used to train in
eager PyTorch Autograd. This API is intended for eager training use cases such as PT2 Quantization
and will soon be the default IR of torch.export.export. To read further about
the motivation behind this change, please refer to
<a class="reference external" href="https://dev-discuss.pytorch.org/t/why-pytorch-does-not-need-a-new-standardized-operator-set/2206">https://dev-discuss.pytorch.org/t/why-pytorch-does-not-need-a-new-standardized-operator-set/2206</a></p>
<p>When this API is combined with <code class="xref py py-func docutils literal notranslate"><span class="pre">run_decompositions()</span></code>, you should be able to get inference IR with
any desired decomposition behavior.</p>
<p>To show some examples:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">ConvBatchnorm</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</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="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">conv</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Conv2d</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">bn</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">BatchNorm2d</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">conv</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">bn</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">x</span><span class="p">,)</span>

<span class="n">mod</span> <span class="o">=</span> <span class="n">ConvBatchnorm</span><span class="p">()</span>
<span class="n">inp</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>

<span class="n">ep_for_training</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export_for_training</span><span class="p">(</span><span class="n">mod</span><span class="p">,</span> <span class="p">(</span><span class="n">inp</span><span class="p">,))</span>
<span class="nb">print</span><span class="p">(</span><span class="n">ep_for_training</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
    <span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
        <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">:</span> <span class="s2">&quot;f32[3, 1, 1, 1]&quot;</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">p_bn_weight</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">p_bn_bias</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_running_mean</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_running_var</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_num_batches_tracked</span><span class="p">:</span> <span class="s2">&quot;i64[]&quot;</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[1, 1, 3, 3]&quot;</span><span class="p">):</span>
            <span class="n">conv2d</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">conv2d</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">)</span>
            <span class="n">add_</span><span class="p">:</span> <span class="s2">&quot;i64[]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add_</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">b_bn_num_batches_tracked</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">batch_norm</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">batch_norm</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">conv2d</span><span class="p">,</span> <span class="n">p_bn_weight</span><span class="p">,</span> <span class="n">p_bn_bias</span><span class="p">,</span> <span class="n">b_bn_running_mean</span><span class="p">,</span> <span class="n">b_bn_running_var</span><span class="p">,</span> <span class="kc">True</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mf">1e-05</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">batch_norm</span><span class="p">,)</span>

<span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
    <span class="n">ExportGraphSignature</span><span class="p">(</span>
        <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_bn_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_bn_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_running_mean&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_mean&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_running_var&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_var&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_num_batches_tracked&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.num_batches_tracked&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">],</span>
        <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;batch_norm&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">]</span>
    <span class="p">)</span>
<span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{}</span>
</pre></div>
</div>
<p>From the above output, you can see that <code class="xref py py-func docutils literal notranslate"><span class="pre">export_for_training()</span></code> produces pretty much the same ExportedProgram
as <code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code> except for the operators in the graph. You can see that we captured batch_norm in the most general
form. This op is non-functional and will be lowered to different ops when running inference.</p>
<p>You can also go from this IR to an inference IR via <code class="xref py py-func docutils literal notranslate"><span class="pre">run_decompositions()</span></code> with arbitrary customizations.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Lower to core aten inference IR, but keep conv2d</span>
<span class="n">decomp_table</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">default_decompositions</span><span class="p">()</span>
<span class="k">del</span> <span class="n">decomp_table</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">conv2d</span><span class="o">.</span><span class="n">default</span><span class="p">]</span>
<span class="n">ep_for_inference</span> <span class="o">=</span> <span class="n">ep_for_training</span><span class="o">.</span><span class="n">run_decompositions</span><span class="p">(</span><span class="n">decomp_table</span><span class="p">)</span>

<span class="nb">print</span><span class="p">(</span><span class="n">ep_for_inference</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
    <span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
        <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">:</span> <span class="s2">&quot;f32[3, 1, 1, 1]&quot;</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">p_bn_weight</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">p_bn_bias</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_running_mean</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_running_var</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_num_batches_tracked</span><span class="p">:</span> <span class="s2">&quot;i64[]&quot;</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[1, 1, 3, 3]&quot;</span><span class="p">):</span>
            <span class="n">conv2d</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">conv2d</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">)</span>
            <span class="n">add</span><span class="p">:</span> <span class="s2">&quot;i64[]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">b_bn_num_batches_tracked</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">_native_batch_norm_legit_functional</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">_native_batch_norm_legit_functional</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">conv2d</span><span class="p">,</span> <span class="n">p_bn_weight</span><span class="p">,</span> <span class="n">p_bn_bias</span><span class="p">,</span> <span class="n">b_bn_running_mean</span><span class="p">,</span> <span class="n">b_bn_running_var</span><span class="p">,</span> <span class="kc">True</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mf">1e-05</span><span class="p">)</span>
            <span class="n">getitem</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">_native_batch_norm_legit_functional</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">getitem_3</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span> <span class="o">=</span> <span class="n">_native_batch_norm_legit_functional</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
            <span class="n">getitem_4</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span> <span class="o">=</span> <span class="n">_native_batch_norm_legit_functional</span><span class="p">[</span><span class="mi">4</span><span class="p">]</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">getitem_3</span><span class="p">,</span> <span class="n">getitem_4</span><span class="p">,</span> <span class="n">add</span><span class="p">,</span> <span class="n">getitem</span><span class="p">)</span>

<span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
    <span class="n">ExportGraphSignature</span><span class="p">(</span>
        <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_bn_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_bn_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_running_mean&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_mean&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_running_var&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_var&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_num_batches_tracked&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.num_batches_tracked&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">],</span>
        <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;getitem_3&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_mean&#39;</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;getitem_4&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_var&#39;</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.num_batches_tracked&#39;</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;getitem&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">]</span>
    <span class="p">)</span>
<span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{}</span>
</pre></div>
</div>
<p>Here you can see that we kept <code class="docutils literal notranslate"><span class="pre">conv2d</span></code> op in the IR while decomposing the rest. Now the IR is a functional IR
containing core aten operators except for <code class="docutils literal notranslate"><span class="pre">conv2d</span></code>.</p>
<p>You can do even more customization by directly registering your chosen decomposition behaviors.</p>
<p>You can do even more customizations by directly registering custom decomp behaviour</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Lower to core aten inference IR, but customize conv2d</span>
<span class="n">decomp_table</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">default_decompositions</span><span class="p">()</span>

<span class="k">def</span> <span class="nf">my_awesome_custom_conv2d_function</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,</span> <span class="n">stride</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">padding</span><span class="o">=</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="n">dilation</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">groups</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
    <span class="k">return</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">convolution</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="kc">False</span><span class="p">,</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="n">groups</span><span class="p">)</span>

<span class="n">decomp_table</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">conv2d</span><span class="o">.</span><span class="n">default</span><span class="p">]</span> <span class="o">=</span> <span class="n">my_awesome_conv2d_function</span>
<span class="n">ep_for_inference</span> <span class="o">=</span> <span class="n">ep_for_training</span><span class="o">.</span><span class="n">run_decompositions</span><span class="p">(</span><span class="n">decomp_table</span><span class="p">)</span>

<span class="nb">print</span><span class="p">(</span><span class="n">ep_for_inference</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
    <span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
        <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">:</span> <span class="s2">&quot;f32[3, 1, 1, 1]&quot;</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">p_bn_weight</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">p_bn_bias</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_running_mean</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_running_var</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span><span class="p">,</span> <span class="n">b_bn_num_batches_tracked</span><span class="p">:</span> <span class="s2">&quot;i64[]&quot;</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[1, 1, 3, 3]&quot;</span><span class="p">):</span>
            <span class="n">convolution</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">convolution</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">p_conv_weight</span><span class="p">,</span> <span class="n">p_conv_bias</span><span class="p">,</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</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="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="kc">False</span><span class="p">,</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="mi">1</span><span class="p">)</span>
            <span class="n">mul</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">convolution</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
            <span class="n">add</span><span class="p">:</span> <span class="s2">&quot;i64[]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">b_bn_num_batches_tracked</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">_native_batch_norm_legit_functional</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">_native_batch_norm_legit_functional</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">mul</span><span class="p">,</span> <span class="n">p_bn_weight</span><span class="p">,</span> <span class="n">p_bn_bias</span><span class="p">,</span> <span class="n">b_bn_running_mean</span><span class="p">,</span> <span class="n">b_bn_running_var</span><span class="p">,</span> <span class="kc">True</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mf">1e-05</span><span class="p">)</span>
            <span class="n">getitem</span><span class="p">:</span> <span class="s2">&quot;f32[1, 3, 3, 3]&quot;</span> <span class="o">=</span> <span class="n">_native_batch_norm_legit_functional</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">getitem_3</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span> <span class="o">=</span> <span class="n">_native_batch_norm_legit_functional</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
            <span class="n">getitem_4</span><span class="p">:</span> <span class="s2">&quot;f32[3]&quot;</span> <span class="o">=</span> <span class="n">_native_batch_norm_legit_functional</span><span class="p">[</span><span class="mi">4</span><span class="p">];</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">getitem_3</span><span class="p">,</span> <span class="n">getitem_4</span><span class="p">,</span> <span class="n">add</span><span class="p">,</span> <span class="n">getitem</span><span class="p">)</span>

<span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
    <span class="n">ExportGraphSignature</span><span class="p">(</span>
        <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_conv_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;conv.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_bn_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_bn_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_running_mean&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_mean&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_running_var&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_var&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;b_bn_num_batches_tracked&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.num_batches_tracked&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">],</span>
        <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;getitem_3&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_mean&#39;</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;getitem_4&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.running_var&#39;</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;bn.num_batches_tracked&#39;</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;getitem&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">]</span>
<span class="p">)</span>
<span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{}</span>
</pre></div>
</div>
</div>
<div class="section" id="expressing-dynamism">
<h3>Expressing Dynamism<a class="headerlink" href="#expressing-dynamism" title="Permalink to this heading">¶</a></h3>
<p>By default <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> will trace the program assuming all input shapes are
<strong>static</strong>, and specializing the exported program to those dimensions. However,
some dimensions, such as a batch dimension, can be dynamic and vary from run to
run. Such dimensions must be specified by using the
<code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.Dim()</span></code> API to create them and by passing them into
<a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.export()</span></code></a> through the <code class="docutils literal notranslate"><span class="pre">dynamic_shapes</span></code> argument. An example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.export</span> <span class="kn">import</span> <span class="n">Dim</span><span class="p">,</span> <span class="n">export</span>

<span class="k">class</span> <span class="nc">M</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</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="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">branch1</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">64</span><span class="p">,</span> <span class="mi">32</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ReLU</span><span class="p">()</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">branch2</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Sequential</span><span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">64</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ReLU</span><span class="p">()</span>
        <span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">buffer</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">32</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">):</span>
        <span class="n">out1</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">branch1</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
        <span class="n">out2</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">branch2</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">out1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">buffer</span><span class="p">,</span> <span class="n">out2</span><span class="p">)</span>

<span class="n">example_args</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">32</span><span class="p">,</span> <span class="mi">64</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">32</span><span class="p">,</span> <span class="mi">128</span><span class="p">))</span>

<span class="c1"># Create a dynamic batch size</span>
<span class="n">batch</span> <span class="o">=</span> <span class="n">Dim</span><span class="p">(</span><span class="s2">&quot;batch&quot;</span><span class="p">)</span>
<span class="c1"># Specify that the first dimension of each input is that batch size</span>
<span class="n">dynamic_shapes</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;x1&quot;</span><span class="p">:</span> <span class="p">{</span><span class="mi">0</span><span class="p">:</span> <span class="n">batch</span><span class="p">},</span> <span class="s2">&quot;x2&quot;</span><span class="p">:</span> <span class="p">{</span><span class="mi">0</span><span class="p">:</span> <span class="n">batch</span><span class="p">}}</span>

<span class="n">exported_program</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">ExportedProgram</span> <span class="o">=</span> <span class="n">export</span><span class="p">(</span>
    <span class="n">M</span><span class="p">(),</span> <span class="n">args</span><span class="o">=</span><span class="n">example_args</span><span class="p">,</span> <span class="n">dynamic_shapes</span><span class="o">=</span><span class="n">dynamic_shapes</span>
<span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">exported_program</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
<span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">p_branch1_0_weight</span><span class="p">:</span> <span class="s2">&quot;f32[32, 64]&quot;</span><span class="p">,</span> <span class="n">p_branch1_0_bias</span><span class="p">:</span> <span class="s2">&quot;f32[32]&quot;</span><span class="p">,</span> <span class="n">p_branch2_0_weight</span><span class="p">:</span> <span class="s2">&quot;f32[64, 128]&quot;</span><span class="p">,</span> <span class="n">p_branch2_0_bias</span><span class="p">:</span> <span class="s2">&quot;f32[64]&quot;</span><span class="p">,</span> <span class="n">c_buffer</span><span class="p">:</span> <span class="s2">&quot;f32[32]&quot;</span><span class="p">,</span> <span class="n">x1</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 64]&quot;</span><span class="p">,</span> <span class="n">x2</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 128]&quot;</span><span class="p">):</span>

         <span class="c1"># code: out1 = self.branch1(x1)</span>
        <span class="n">linear</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 32]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">linear</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">p_branch1_0_weight</span><span class="p">,</span> <span class="n">p_branch1_0_bias</span><span class="p">)</span>
        <span class="n">relu</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 32]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">relu</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">linear</span><span class="p">)</span>

         <span class="c1"># code: out2 = self.branch2(x2)</span>
        <span class="n">linear_1</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 64]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">linear</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="n">p_branch2_0_weight</span><span class="p">,</span> <span class="n">p_branch2_0_bias</span><span class="p">)</span>
        <span class="n">relu_1</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 64]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">relu</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="n">linear_1</span><span class="p">)</span>

         <span class="c1"># code: return (out1 + self.buffer, out2)</span>
        <span class="n">add</span><span class="p">:</span> <span class="s2">&quot;f32[s0, 32]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">relu</span><span class="p">,</span> <span class="n">c_buffer</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">add</span><span class="p">,</span> <span class="n">relu_1</span><span class="p">)</span>

<span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
    <span class="n">ExportGraphSignature</span><span class="p">(</span>
        <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_branch1_0_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;branch1.0.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_branch1_0_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;branch1.0.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_branch2_0_weight&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;branch2.0.weight&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;p_branch2_0_bias&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;branch2.0.bias&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">CONSTANT_TENSOR</span><span class="p">:</span> <span class="mi">4</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;c_buffer&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="s1">&#39;buffer&#39;</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x1&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x2&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">],</span>
        <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;relu_1&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">]</span>
    <span class="p">)</span>
<span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{</span><span class="n">s0</span><span class="p">:</span> <span class="n">VR</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="n">int_oo</span><span class="p">]}</span>
</pre></div>
</div>
<p>Some additional things to note:</p>
<ul class="simple">
<li><p>Through the <code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.Dim()</span></code> API and the <code class="docutils literal notranslate"><span class="pre">dynamic_shapes</span></code> argument, we specified the first
dimension of each input to be dynamic. Looking at the inputs <code class="docutils literal notranslate"><span class="pre">x1</span></code> and
<code class="docutils literal notranslate"><span class="pre">x2</span></code>, they have a symbolic shape of (s0, 64) and (s0, 128), instead of
the (32, 64) and (32, 128) shaped tensors that we passed in as example inputs.
<code class="docutils literal notranslate"><span class="pre">s0</span></code> is a symbol representing that this dimension can be a range
of values.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">exported_program.range_constraints</span></code> describes the ranges of each symbol
appearing in the graph. In this case, we see that <code class="docutils literal notranslate"><span class="pre">s0</span></code> has the range
[0, int_oo]. For technical reasons that are difficult to explain here, they are
assumed to be not 0 or 1. This is not a bug, and does not necessarily mean
that the exported program will not work for dimensions 0 or 1. See
<a class="reference external" href="https://docs.google.com/document/d/16VPOa3d-Liikf48teAOmxLc92rgvJdfosIy-yoT38Io/edit?fbclid=IwAR3HNwmmexcitV0pbZm_x1a4ykdXZ9th_eJWK-3hBtVgKnrkmemz6Pm5jRQ#heading=h.ez923tomjvyk">The 0/1 Specialization Problem</a>
for an in-depth discussion of this topic.</p></li>
</ul>
<p>We can also specify more expressive relationships between input shapes, such as
where a pair of shapes might differ by one, a shape might be double of
another, or a shape is even. An example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">M</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x</span> <span class="o">+</span> <span class="n">y</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>

<span class="n">x</span><span class="p">,</span> <span class="n">y</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">5</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">6</span><span class="p">)</span>
<span class="n">dimx</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">Dim</span><span class="p">(</span><span class="s2">&quot;dimx&quot;</span><span class="p">,</span> <span class="nb">min</span><span class="o">=</span><span class="mi">3</span><span class="p">,</span> <span class="nb">max</span><span class="o">=</span><span class="mi">6</span><span class="p">)</span>
<span class="n">dimy</span> <span class="o">=</span> <span class="n">dimx</span> <span class="o">+</span> <span class="mi">1</span>

<span class="n">exported_program</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export</span><span class="p">(</span>
    <span class="n">M</span><span class="p">(),</span> <span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">),</span> <span class="n">dynamic_shapes</span><span class="o">=</span><span class="p">({</span><span class="mi">0</span><span class="p">:</span> <span class="n">dimx</span><span class="p">},</span> <span class="p">{</span><span class="mi">0</span><span class="p">:</span> <span class="n">dimy</span><span class="p">}),</span>
<span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">exported_program</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
<span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[s0]&quot;</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="s2">&quot;f32[s0 + 1]&quot;</span><span class="p">):</span>
        <span class="c1"># code: return x + y[1:]</span>
        <span class="n">slice_1</span><span class="p">:</span> <span class="s2">&quot;f32[s0]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">slice</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">9223372036854775807</span><span class="p">)</span>
        <span class="n">add</span><span class="p">:</span> <span class="s2">&quot;f32[s0]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">slice_1</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">add</span><span class="p">,)</span>

<span class="n">Graph</span> <span class="n">signature</span><span class="p">:</span>
    <span class="n">ExportGraphSignature</span><span class="p">(</span>
        <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;x&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">),</span>
            <span class="n">InputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;y&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                <span class="n">persistent</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">],</span>
        <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
            <span class="n">OutputSpec</span><span class="p">(</span>
                <span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span>
                <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add&#39;</span><span class="p">),</span>
                <span class="n">target</span><span class="o">=</span><span class="kc">None</span>
            <span class="p">)</span>
        <span class="p">]</span>
    <span class="p">)</span>
<span class="n">Range</span> <span class="n">constraints</span><span class="p">:</span> <span class="p">{</span><span class="n">s0</span><span class="p">:</span> <span class="n">VR</span><span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">6</span><span class="p">],</span> <span class="n">s0</span> <span class="o">+</span> <span class="mi">1</span><span class="p">:</span> <span class="n">VR</span><span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">7</span><span class="p">]}</span>
</pre></div>
</div>
<p>Some things to note:</p>
<ul class="simple">
<li><p>By specifying <code class="docutils literal notranslate"><span class="pre">{0:</span> <span class="pre">dimx}</span></code> for the first input, we see that the resulting
shape of the first input is now dynamic, being <code class="docutils literal notranslate"><span class="pre">[s0]</span></code>. And now by specifying
<code class="docutils literal notranslate"><span class="pre">{0:</span> <span class="pre">dimy}</span></code> for the second input, we see that the resulting shape of the
second input is also dynamic. However, because we expressed <code class="docutils literal notranslate"><span class="pre">dimy</span> <span class="pre">=</span> <span class="pre">dimx</span> <span class="pre">+</span> <span class="pre">1</span></code>,
instead of <code class="docutils literal notranslate"><span class="pre">y</span></code>’s shape containing a new symbol, we see that it is
now being represented with the same symbol used in <code class="docutils literal notranslate"><span class="pre">x</span></code>, <code class="docutils literal notranslate"><span class="pre">s0</span></code>. We can
see that relationship of <code class="docutils literal notranslate"><span class="pre">dimy</span> <span class="pre">=</span> <span class="pre">dimx</span> <span class="pre">+</span> <span class="pre">1</span></code> is being shown through <code class="docutils literal notranslate"><span class="pre">s0</span> <span class="pre">+</span> <span class="pre">1</span></code>.</p></li>
<li><p>Looking at the range constraints, we see that <code class="docutils literal notranslate"><span class="pre">s0</span></code> has the range [3, 6],
which is specified initially, and we can see that <code class="docutils literal notranslate"><span class="pre">s0</span> <span class="pre">+</span> <span class="pre">1</span></code> has the solved
range of [4, 7].</p></li>
</ul>
</div>
<div class="section" id="serialization">
<h3>Serialization<a class="headerlink" href="#serialization" title="Permalink to this heading">¶</a></h3>
<p>To save the <code class="docutils literal notranslate"><span class="pre">ExportedProgram</span></code>, users can use the <a class="reference internal" href="#torch.export.save" title="torch.export.save"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.save()</span></code></a> and
<a class="reference internal" href="#torch.export.load" title="torch.export.load"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.load()</span></code></a> APIs. A convention is to save the <code class="docutils literal notranslate"><span class="pre">ExportedProgram</span></code>
using a <code class="docutils literal notranslate"><span class="pre">.pt2</span></code> file extension.</p>
<p>An example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">io</span>

<span class="k">class</span> <span class="nc">MyModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x</span> <span class="o">+</span> <span class="mi">10</span>

<span class="n">exported_program</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export</span><span class="p">(</span><span class="n">MyModule</span><span class="p">(),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">5</span><span class="p">))</span>

<span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">exported_program</span><span class="p">,</span> <span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">)</span>
<span class="n">saved_exported_program</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="section" id="specializations">
<h3>Specializations<a class="headerlink" href="#specializations" title="Permalink to this heading">¶</a></h3>
<p>A key concept in understanding the behavior of <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> is the
difference between <em>static</em> and <em>dynamic</em> values.</p>
<p>A <em>dynamic</em> value is one that can change from run to run. These behave like
normal arguments to a Python function—you can pass different values for an
argument and expect your function to do the right thing. Tensor <em>data</em> is
treated as dynamic.</p>
<p>A <em>static</em> value is a value that is fixed at export time and cannot change
between executions of the exported program. When the value is encountered during
tracing, the exporter will treat it as a constant and hard-code it into the
graph.</p>
<p>When an operation is performed (e.g. <code class="docutils literal notranslate"><span class="pre">x</span> <span class="pre">+</span> <span class="pre">y</span></code>) and all inputs are static, then
the output of the operation will be directly hard-coded into the graph, and the
operation won’t show up (i.e. it will get constant-folded).</p>
<p>When a value has been hard-coded into the graph, we say that the graph has been
<em>specialized</em> to that value.</p>
<p>The following values are static:</p>
<div class="section" id="input-tensor-shapes">
<h4>Input Tensor Shapes<a class="headerlink" href="#input-tensor-shapes" title="Permalink to this heading">¶</a></h4>
<p>By default, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> will trace the program specializing on the input
tensors’ shapes, unless a dimension is specified as dynamic via the
<code class="docutils literal notranslate"><span class="pre">dynamic_shapes</span></code> argument to <code class="docutils literal notranslate"><span class="pre">torch.export</span></code>. This means that if there exists
shape-dependent control flow, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> will specialize on the branch
that is being taken with the given sample inputs. For example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.export</span> <span class="kn">import</span> <span class="n">export</span>

<span class="k">class</span> <span class="nc">Mod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">x</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">5</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">x</span> <span class="o">+</span> <span class="mi">1</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">x</span> <span class="o">-</span> <span class="mi">1</span>

<span class="n">example_inputs</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">10</span><span class="p">,</span> <span class="mi">2</span><span class="p">),)</span>
<span class="n">exported_program</span> <span class="o">=</span> <span class="n">export</span><span class="p">(</span><span class="n">Mod</span><span class="p">(),</span> <span class="n">example_inputs</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">exported_program</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
<span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[10, 2]&quot;</span><span class="p">):</span>
        <span class="c1"># code: return x + 1</span>
        <span class="n">add</span><span class="p">:</span> <span class="s2">&quot;f32[10, 2]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">add</span><span class="p">,)</span>
</pre></div>
</div>
<p>The conditional of (<code class="docutils literal notranslate"><span class="pre">x.shape[0]</span> <span class="pre">&gt;</span> <span class="pre">5</span></code>) does not appear in the
<code class="docutils literal notranslate"><span class="pre">ExportedProgram</span></code> because the example inputs have the static
shape of (10, 2). Since <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> specializes on the inputs’ static
shapes, the else branch (<code class="docutils literal notranslate"><span class="pre">x</span> <span class="pre">-</span> <span class="pre">1</span></code>) will never be reached. To preserve the dynamic
branching behavior based on the shape of a tensor in the traced graph,
<code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.Dim()</span></code> will need to be used to specify the dimension
of the input tensor (<code class="docutils literal notranslate"><span class="pre">x.shape[0]</span></code>) to be dynamic, and the source code will
need to be <a class="reference internal" href="#data-shape-dependent-control-flow"><span class="std std-ref">rewritten</span></a>.</p>
<p>Note that tensors that are part of the module state (e.g. parameters and
buffers) always have static shapes.</p>
</div>
<div class="section" id="python-primitives">
<h4>Python Primitives<a class="headerlink" href="#python-primitives" title="Permalink to this heading">¶</a></h4>
<p><code class="docutils literal notranslate"><span class="pre">torch.export</span></code> also specializes on Python primtivies,
such as <code class="docutils literal notranslate"><span class="pre">int</span></code>, <code class="docutils literal notranslate"><span class="pre">float</span></code>, <code class="docutils literal notranslate"><span class="pre">bool</span></code>, and <code class="docutils literal notranslate"><span class="pre">str</span></code>. However they do have dynamic
variants such as <code class="docutils literal notranslate"><span class="pre">SymInt</span></code>, <code class="docutils literal notranslate"><span class="pre">SymFloat</span></code>, and <code class="docutils literal notranslate"><span class="pre">SymBool</span></code>.</p>
<p>For example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch.export</span> <span class="kn">import</span> <span class="n">export</span>

<span class="k">class</span> <span class="nc">Mod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</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">const</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">times</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">times</span><span class="p">):</span>
            <span class="n">x</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="n">const</span>
        <span class="k">return</span> <span class="n">x</span>

<span class="n">example_inputs</span> <span class="o">=</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
<span class="n">exported_program</span> <span class="o">=</span> <span class="n">export</span><span class="p">(</span><span class="n">Mod</span><span class="p">(),</span> <span class="n">example_inputs</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">exported_program</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportedProgram</span><span class="p">:</span>
    <span class="k">class</span> <span class="nc">GraphModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
        <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="s2">&quot;f32[2, 2]&quot;</span><span class="p">,</span> <span class="n">const</span><span class="p">,</span> <span class="n">times</span><span class="p">):</span>
            <span class="c1"># code: x = x + const</span>
            <span class="n">add</span><span class="p">:</span> <span class="s2">&quot;f32[2, 2]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">add_1</span><span class="p">:</span> <span class="s2">&quot;f32[2, 2]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">add</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">add_2</span><span class="p">:</span> <span class="s2">&quot;f32[2, 2]&quot;</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">(</span><span class="n">add_1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">add_2</span><span class="p">,)</span>
</pre></div>
</div>
<p>Because integers are specialized, the <code class="docutils literal notranslate"><span class="pre">torch.ops.aten.add.Tensor</span></code> operations
are all computed with the hard-coded constant <code class="docutils literal notranslate"><span class="pre">1</span></code>, rather than <code class="docutils literal notranslate"><span class="pre">const</span></code>. If
a user passes a different value for <code class="docutils literal notranslate"><span class="pre">const</span></code> at runtime, like 2, than the one used
during export time, 1, this will result in an error.
Additionally, the <code class="docutils literal notranslate"><span class="pre">times</span></code> iterator used in the <code class="docutils literal notranslate"><span class="pre">for</span></code> loop is also “inlined”
in the graph through the 3 repeated <code class="docutils literal notranslate"><span class="pre">torch.ops.aten.add.Tensor</span></code> calls, and the
input <code class="docutils literal notranslate"><span class="pre">times</span></code> is never used.</p>
</div>
<div class="section" id="python-containers">
<h4>Python Containers<a class="headerlink" href="#python-containers" title="Permalink to this heading">¶</a></h4>
<p>Python containers (<code class="docutils literal notranslate"><span class="pre">List</span></code>, <code class="docutils literal notranslate"><span class="pre">Dict</span></code>, <code class="docutils literal notranslate"><span class="pre">NamedTuple</span></code>, etc.) are considered to
have static structure.</p>
</div>
</div>
</div>
<div class="section" id="limitations-of-torch-export">
<span id="id3"></span><h2>Limitations of torch.export<a class="headerlink" href="#limitations-of-torch-export" title="Permalink to this heading">¶</a></h2>
<div class="section" id="graph-breaks">
<h3>Graph Breaks<a class="headerlink" href="#graph-breaks" title="Permalink to this heading">¶</a></h3>
<p>As <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> is a one-shot process for capturing a computation graph from
a PyTorch program, it might ultimately run into untraceable parts of programs as
it is nearly impossible to support tracing all PyTorch and Python features. In
the case of <code class="docutils literal notranslate"><span class="pre">torch.compile</span></code>, an unsupported operation will cause a “graph
break” and the unsupported operation will be run with default Python evaluation.
In contrast, <code class="docutils literal notranslate"><span class="pre">torch.export</span></code> will require users to provide additional
information or rewrite parts of their code to make it traceable. As the
tracing is based on TorchDynamo, which evaluates at the Python
bytecode level, there will be significantly fewer rewrites required compared to
previous tracing frameworks.</p>
<p>When a graph break is encountered, <a class="reference internal" href="generated/exportdb/index.html#torch-export-db"><span class="std std-ref">ExportDB</span></a> is a great
resource for learning about the kinds of programs that are supported and
unsupported, along with ways to rewrite programs to make them traceable.</p>
<p>An option to get past dealing with this graph breaks is by using
<a class="reference internal" href="#non-strict-export"><span class="std std-ref">non-strict export</span></a></p>
</div>
<div class="section" id="data-shape-dependent-control-flow">
<span id="id4"></span><h3>Data/Shape-Dependent Control Flow<a class="headerlink" href="#data-shape-dependent-control-flow" title="Permalink to this heading">¶</a></h3>
<p>Graph breaks can also be encountered on data-dependent control flow (<code class="docutils literal notranslate"><span class="pre">if</span>
<span class="pre">x.shape[0]</span> <span class="pre">&gt;</span> <span class="pre">2</span></code>) when shapes are not being specialized, as a tracing compiler cannot
possibly deal with without generating code for a combinatorially exploding
number of paths. In such cases, users will need to rewrite their code using
special control flow operators. Currently, we support <a class="reference internal" href="cond.html#cond"><span class="std std-ref">torch.cond</span></a>
to express if-else like control flow (more coming soon!).</p>
</div>
<div class="section" id="missing-fake-meta-abstract-kernels-for-operators">
<h3>Missing Fake/Meta/Abstract Kernels for Operators<a class="headerlink" href="#missing-fake-meta-abstract-kernels-for-operators" title="Permalink to this heading">¶</a></h3>
<p>When tracing, a FakeTensor kernel (aka meta kernel, abstract impl) is
required for all operators. This is used to reason about the input/output shapes
for this operator.</p>
<p>Please see <a class="reference internal" href="library.html#torch.library.register_fake" title="torch.library.register_fake"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.library.register_fake()</span></code></a> for more details.</p>
<p>In the unfortunate case where your model uses an ATen operator that is does not
have a FakeTensor kernel implementation yet, please file an issue.</p>
</div>
</div>
<div class="section" id="read-more">
<h2>Read More<a class="headerlink" href="#read-more" title="Permalink to this heading">¶</a></h2>
<div class="toctree-wrapper compound">
<p class="caption" role="heading"><span class="caption-text">Additional Links for Export Users</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="export.ir_spec.html">torch.export IR Specification</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler_transformations.html">Writing Graph Transformations on ATen IR</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler_ir.html">IRs</a></li>
<li class="toctree-l1"><a class="reference internal" href="generated/exportdb/index.html">ExportDB</a></li>
<li class="toctree-l1"><a class="reference internal" href="cond.html">Control Flow - Cond</a></li>
</ul>
</div>
<div class="toctree-wrapper compound">
<p class="caption" role="heading"><span class="caption-text">Deep Dive for PyTorch Developers</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler_dynamo_overview.html">Dynamo Overview</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler_dynamo_deepdive.html">Dynamo Deep-Dive</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler_dynamic_shapes.html">Dynamic shapes</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.compiler_fake_tensor.html">Fake tensor</a></li>
</ul>
</div>
</div>
<div class="section" id="module-torch.export">
<span id="api-reference"></span><h2>API Reference<a class="headerlink" href="#module-torch.export" title="Permalink to this heading">¶</a></h2>
<dl class="py function">
<dt class="sig sig-object py" id="torch.export.export">
<span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">export</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">mod</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">args</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">kwargs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dynamic_shapes</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">strict</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">preserve_module_call_signature</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">()</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export.html#export"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/__init__.py#L263"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.export" title="Permalink to this definition">¶</a></dt>
<dd><p><a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a> takes any nn.Module along with example inputs, and produces a traced graph representing
only the Tensor computation of the function in an Ahead-of-Time (AOT) fashion,
which can subsequently be executed with different inputs or serialized.  The
traced graph (1) produces normalized operators in the functional ATen operator set
(as well as any user-specified custom operators), (2) has eliminated all Python control
flow and data structures (with certain exceptions), and (3) records the set of
shape constraints needed to show that this normalization and control-flow elimination
is sound for future inputs.</p>
<p><strong>Soundness Guarantee</strong></p>
<p>While tracing, <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a> takes note of shape-related assumptions
made by the user program and the underlying PyTorch operator kernels.
The output <a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportedProgram</span></code></a> is considered valid only when these
assumptions hold true.</p>
<p>Tracing makes assumptions on the shapes (not values) of input tensors.
Such assumptions must be validated at graph capture time for <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a>
to succeed. Specifically:</p>
<ul class="simple">
<li><p>Assumptions on static shapes of input tensors are automatically validated without additional effort.</p></li>
<li><p>Assumptions on dynamic shape of input tensors require explicit specification
by using the <code class="xref py py-func docutils literal notranslate"><span class="pre">Dim()</span></code> API to construct dynamic dimensions and by associating
them with example inputs through the <code class="docutils literal notranslate"><span class="pre">dynamic_shapes</span></code> argument.</p></li>
</ul>
<p>If any assumption can not be validated, a fatal error will be raised. When that happens,
the error message will include suggested fixes to the specification that are needed
to validate the assumptions. For example <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a> might suggest the
following fix to the definition of a dynamic dimension <code class="docutils literal notranslate"><span class="pre">dim0_x</span></code>, say appearing in the
shape associated with input <code class="docutils literal notranslate"><span class="pre">x</span></code>, that was previously defined as <code class="docutils literal notranslate"><span class="pre">Dim(&quot;dim0_x&quot;)</span></code>:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">dim</span> <span class="o">=</span> <span class="n">Dim</span><span class="p">(</span><span class="s2">&quot;dim0_x&quot;</span><span class="p">,</span> <span class="nb">max</span><span class="o">=</span><span class="mi">5</span><span class="p">)</span>
</pre></div>
</div>
<p>This example means the generated code requires dimension 0 of input <code class="docutils literal notranslate"><span class="pre">x</span></code> to be less
than or equal to 5 to be valid. You can inspect the suggested fixes to dynamic dimension
definitions and then copy them verbatim into your code without needing to change the
<code class="docutils literal notranslate"><span class="pre">dynamic_shapes</span></code> argument to your <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a> call.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>mod</strong> (<a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.modules.module.Module"><em>Module</em></a>) – We will trace the forward method of this module.</p></li>
<li><p><strong>args</strong> (<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.13)"><em>Tuple</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a><em>, </em><em>...</em><em>]</em>) – Example positional inputs.</p></li>
<li><p><strong>kwargs</strong> (<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Optional" title="(in Python v3.13)"><em>Optional</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.13)"><em>Dict</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a><em>]</em><em>]</em>) – Optional example keyword inputs.</p></li>
<li><p><strong>dynamic_shapes</strong> (<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Optional" title="(in Python v3.13)"><em>Optional</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Union" title="(in Python v3.13)"><em>Union</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.13)"><em>Dict</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a><em>]</em><em>, </em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.13)"><em>Tuple</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a><em>]</em><em>, </em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.List" title="(in Python v3.13)"><em>List</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a><em>]</em><em>]</em><em>]</em>) – <p>An optional argument where the type should either be:
1) a dict from argument names of <code class="docutils literal notranslate"><span class="pre">f</span></code> to their dynamic shape specifications,
2) a tuple that specifies dynamic shape specifications for each input in original order.
If you are specifying dynamism on keyword args, you will need to pass them in the order that
is defined in the original function signature.</p>
<p>The dynamic shape of a tensor argument can be specified as either
(1) a dict from dynamic dimension indices to <code class="xref py py-func docutils literal notranslate"><span class="pre">Dim()</span></code> types, where it is
not required to include static dimension indices in this dict, but when they are,
they should be mapped to None; or (2) a tuple / list of <code class="xref py py-func docutils literal notranslate"><span class="pre">Dim()</span></code> types or None,
where the <code class="xref py py-func docutils literal notranslate"><span class="pre">Dim()</span></code> types correspond to dynamic dimensions, and static dimensions
are denoted by None. Arguments that are dicts or tuples / lists of tensors are
recursively specified by using mappings or sequences of contained specifications.</p>
</p></li>
<li><p><strong>strict</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)"><em>bool</em></a>) – When enabled (default), the export function will trace the program through
TorchDynamo which will ensure the soundness of the resulting graph. Otherwise, the
exported program will not validate the implicit assumptions baked into the graph and
may cause behavior divergence between the original model and the exported one. This is
useful when users need to workaround bugs in the tracer, or simply want incrementally
enable safety in their models. Note that this does not affect the resulting IR spec
to be different and the model will be serialized in the same way regardless of what value
is passed here.
WARNING: This option is experimental and use this at your own risk.</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>An <a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportedProgram</span></code></a> containing the traced callable.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.exported_program.ExportedProgram"><em>ExportedProgram</em></a></p>
</dd>
</dl>
<p><strong>Acceptable input/output types</strong></p>
<p>Acceptable types of inputs (for <code class="docutils literal notranslate"><span class="pre">args</span></code> and <code class="docutils literal notranslate"><span class="pre">kwargs</span></code>) and outputs include:</p>
<ul class="simple">
<li><p>Primitive types, i.e. <code class="docutils literal notranslate"><span class="pre">torch.Tensor</span></code>, <code class="docutils literal notranslate"><span class="pre">int</span></code>, <code class="docutils literal notranslate"><span class="pre">float</span></code>, <code class="docutils literal notranslate"><span class="pre">bool</span></code> and <code class="docutils literal notranslate"><span class="pre">str</span></code>.</p></li>
<li><p>Dataclasses, but they must be registered by calling <a class="reference internal" href="#torch.export.register_dataclass" title="torch.export.register_dataclass"><code class="xref py py-func docutils literal notranslate"><span class="pre">register_dataclass()</span></code></a> first.</p></li>
<li><p>(Nested) Data structures comprising of <code class="docutils literal notranslate"><span class="pre">dict</span></code>, <code class="docutils literal notranslate"><span class="pre">list</span></code>, <code class="docutils literal notranslate"><span class="pre">tuple</span></code>, <code class="docutils literal notranslate"><span class="pre">namedtuple</span></code> and
<code class="docutils literal notranslate"><span class="pre">OrderedDict</span></code> containing all above types.</p></li>
</ul>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.save">
<span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">save</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">ep</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">f</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">extra_files</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">opset_version</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export.html#save"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/__init__.py#L379"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.save" title="Permalink to this definition">¶</a></dt>
<dd><div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Under active development, saved files may not be usable in newer versions
of PyTorch.</p>
</div>
<p>Saves an <a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportedProgram</span></code></a> to a file-like object. It can then be
loaded using the Python API <a class="reference internal" href="#torch.export.load" title="torch.export.load"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.load</span></code></a>.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>ep</strong> (<a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><em>ExportedProgram</em></a>) – The exported program to save.</p></li>
<li><p><strong>f</strong> (<em>Union</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/os.html#os.PathLike" title="(in Python v3.13)"><em>os.PathLike</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/io.html#io.BytesIO" title="(in Python v3.13)"><em>io.BytesIO</em></a>) – A file-like object (has to
implement write and flush) or a string containing a file name.</p></li>
<li><p><strong>extra_files</strong> (<em>Optional</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><em>Any</em><em>]</em><em>]</em>) – Map from filename to contents
which will be stored as part of f.</p></li>
<li><p><strong>opset_version</strong> (<em>Optional</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)"><em>int</em></a><em>]</em><em>]</em>) – A map of opset names
to the version of this opset</p></li>
</ul>
</dd>
</dl>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">io</span>

<span class="k">class</span> <span class="nc">MyModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">x</span> <span class="o">+</span> <span class="mi">10</span>

<span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export</span><span class="p">(</span><span class="n">MyModule</span><span class="p">(),</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">5</span><span class="p">),))</span>

<span class="c1"># Save to file</span>
<span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">ep</span><span class="p">,</span> <span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">)</span>

<span class="c1"># Save to io.BytesIO buffer</span>
<span class="n">buffer</span> <span class="o">=</span> <span class="n">io</span><span class="o">.</span><span class="n">BytesIO</span><span class="p">()</span>
<span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">ep</span><span class="p">,</span> <span class="n">buffer</span><span class="p">)</span>

<span class="c1"># Save with extra files</span>
<span class="n">extra_files</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;foo.txt&#39;</span><span class="p">:</span> <span class="sa">b</span><span class="s1">&#39;bar&#39;</span><span class="o">.</span><span class="n">decode</span><span class="p">(</span><span class="s1">&#39;utf-8&#39;</span><span class="p">)}</span>
<span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">ep</span><span class="p">,</span> <span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">,</span> <span class="n">extra_files</span><span class="o">=</span><span class="n">extra_files</span><span class="p">)</span>
</pre></div>
</div>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.load">
<span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">load</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">f</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">extra_files</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">expected_opset_version</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export.html#load"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/__init__.py#L461"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.load" title="Permalink to this definition">¶</a></dt>
<dd><div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Under active development, saved files may not be usable in newer versions
of PyTorch.</p>
</div>
<p>Loads an <a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportedProgram</span></code></a> previously saved with
<a class="reference internal" href="#torch.export.save" title="torch.export.save"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.save</span></code></a>.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>ep</strong> (<a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><em>ExportedProgram</em></a>) – The exported program to save.</p></li>
<li><p><strong>f</strong> (<em>Union</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/os.html#os.PathLike" title="(in Python v3.13)"><em>os.PathLike</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/io.html#io.BytesIO" title="(in Python v3.13)"><em>io.BytesIO</em></a>) – A file-like object (has to
implement write and flush) or a string containing a file name.</p></li>
<li><p><strong>extra_files</strong> (<em>Optional</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><em>Any</em><em>]</em><em>]</em>) – The extra filenames given in
this map would be loaded and their content would be stored in the
provided map.</p></li>
<li><p><strong>expected_opset_version</strong> (<em>Optional</em><em>[</em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)"><em>int</em></a><em>]</em><em>]</em>) – A map of opset names
to expected opset versions</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>An <a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportedProgram</span></code></a> object</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.exported_program.ExportedProgram"><em>ExportedProgram</em></a></p>
</dd>
</dl>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">io</span>

<span class="c1"># Load ExportedProgram from file</span>
<span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">)</span>

<span class="c1"># Load ExportedProgram from io.BytesIO object</span>
<span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">,</span> <span class="s1">&#39;rb&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
    <span class="n">buffer</span> <span class="o">=</span> <span class="n">io</span><span class="o">.</span><span class="n">BytesIO</span><span class="p">(</span><span class="n">f</span><span class="o">.</span><span class="n">read</span><span class="p">())</span>
<span class="n">buffer</span><span class="o">.</span><span class="n">seek</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="n">buffer</span><span class="p">)</span>

<span class="c1"># Load with extra files.</span>
<span class="n">extra_files</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;foo.txt&#39;</span><span class="p">:</span> <span class="s1">&#39;&#39;</span><span class="p">}</span>  <span class="c1"># values will be replaced with data</span>
<span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="s1">&#39;exported_program.pt2&#39;</span><span class="p">,</span> <span class="n">extra_files</span><span class="o">=</span><span class="n">extra_files</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">extra_files</span><span class="p">[</span><span class="s1">&#39;foo.txt&#39;</span><span class="p">])</span>
<span class="nb">print</span><span class="p">(</span><span class="n">ep</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">5</span><span class="p">)))</span>
</pre></div>
</div>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.register_dataclass">
<span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">register_dataclass</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">cls</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">serialized_type_name</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export.html#register_dataclass"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/__init__.py#L576"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.register_dataclass" title="Permalink to this definition">¶</a></dt>
<dd><p>Registers a dataclass as a valid input/output type for <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.export()</span></code></a>.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>cls</strong> (<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Type" title="(in Python v3.13)"><em>Type</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a><em>]</em>) – the dataclass type to register</p></li>
<li><p><strong>serialized_type_name</strong> (<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Optional" title="(in Python v3.13)"><em>Optional</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>]</em>) – The serialized name for the dataclass. This is</p></li>
<li><p><strong>this</strong> (<em>required if you want to serialize the pytree TreeSpec containing</em>) – </p></li>
<li><p><strong>dataclass.</strong> – </p></li>
</ul>
</dd>
</dl>
<p>Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">dataclasses</span> <span class="kn">import</span> <span class="n">dataclass</span>

<span class="nd">@dataclass</span>
<span class="k">class</span> <span class="nc">InputDataClass</span><span class="p">:</span>
    <span class="n">feature</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span>
    <span class="n">bias</span><span class="p">:</span> <span class="nb">int</span>

<span class="nd">@dataclass</span>
<span class="k">class</span> <span class="nc">OutputDataClass</span><span class="p">:</span>
    <span class="n">res</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span>

<span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">register_dataclass</span><span class="p">(</span><span class="n">InputDataClass</span><span class="p">)</span>
<span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">register_dataclass</span><span class="p">(</span><span class="n">OutputDataClass</span><span class="p">)</span>

<span class="k">class</span> <span class="nc">Mod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">:</span> <span class="n">InputDataClass</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">OutputDataClass</span><span class="p">:</span>
        <span class="n">res</span> <span class="o">=</span> <span class="n">x</span><span class="o">.</span><span class="n">feature</span> <span class="o">+</span> <span class="n">x</span><span class="o">.</span><span class="n">bias</span>
        <span class="k">return</span> <span class="n">OutputDataClass</span><span class="p">(</span><span class="n">res</span><span class="o">=</span><span class="n">res</span><span class="p">)</span>

<span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export</span><span class="p">(</span><span class="n">Mod</span><span class="p">(),</span> <span class="p">(</span><span class="n">InputDataClass</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="mi">1</span><span class="p">),</span> <span class="p">))</span>
<span class="nb">print</span><span class="p">(</span><span class="n">ep</span><span class="p">)</span>
</pre></div>
</div>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.dynamic_shapes.Dim">
<span class="sig-prename descclassname"><span class="pre">torch.export.dynamic_shapes.</span></span><span class="sig-name descname"><span class="pre">Dim</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">name</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">min</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">max</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/dynamic_shapes.html#Dim"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/dynamic_shapes.py#L209"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.dynamic_shapes.Dim" title="Permalink to this definition">¶</a></dt>
<dd><p><a class="reference internal" href="#torch.export.dynamic_shapes.Dim" title="torch.export.dynamic_shapes.Dim"><code class="xref py py-func docutils literal notranslate"><span class="pre">Dim()</span></code></a> constructs a type analogous to a named symbolic integer with a range.
It can be used to describe multiple possible values of a dynamic tensor dimension.
Note that different dynamic dimensions of the same tensor, or of different tensors,
can be described by the same type.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>name</strong> (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a>) – Human-readable name for debugging.</p></li>
<li><p><strong>min</strong> (<em>Optional</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)"><em>int</em></a><em>]</em>) – Minimum possible value of given symbol (inclusive)</p></li>
<li><p><strong>max</strong> (<em>Optional</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.13)"><em>int</em></a><em>]</em>) – Maximum possible value of given symbol (inclusive)</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>A type that can be used in dynamic shape specifications for tensors.</p>
</dd>
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.exported_program.default_decompositions">
<span class="sig-prename descclassname"><span class="pre">torch.export.exported_program.</span></span><span class="sig-name descname"><span class="pre">default_decompositions</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#default_decompositions"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L346"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.exported_program.default_decompositions" title="Permalink to this definition">¶</a></dt>
<dd><p>This is the default decomposition table which contains decomposition of
all ATEN operators to core aten opset. Use this API together with
<code class="xref py py-func docutils literal notranslate"><span class="pre">run_decompositions()</span></code></p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable" title="torch.export.decomp_utils.CustomDecompTable"><em>CustomDecompTable</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.dims">
<span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">dims</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="o"><span class="pre">*</span></span><span class="n"><span class="pre">names</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">min</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">max</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/dynamic_shapes.html#dims"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/dynamic_shapes.py#L243"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.dims" title="Permalink to this definition">¶</a></dt>
<dd><p>Util to create multiple <code class="xref py py-func docutils literal notranslate"><span class="pre">Dim()</span></code> types.</p>
<dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.dynamic_shapes.ShapesCollection">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.dynamic_shapes.</span></span><span class="sig-name descname"><span class="pre">ShapesCollection</span></span><a class="reference internal" href="_modules/torch/export/dynamic_shapes.html#ShapesCollection"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/dynamic_shapes.py#L596"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.dynamic_shapes.ShapesCollection" title="Permalink to this definition">¶</a></dt>
<dd><p>Builder for dynamic_shapes.
Used to assign dynamic shape specifications to tensors that appear in inputs.</p>
<dl>
<dt>Example::</dt><dd><p>args = ({“x”: tensor_x, “others”: [tensor_y, tensor_z]})</p>
<p>dim = torch.export.Dim(…)
dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[tensor_y] = {0: dim * 2}
# This is equivalent to the following (now auto-generated):
# dynamic_shapes = {“x”: (dim, dim + 1, 8), “others”: [{0: dim * 2}, None]}</p>
<p>torch.export(…, args, dynamic_shapes=dynamic_shapes)</p>
</dd>
</dl>
<dl class="py method">
<dt class="sig sig-object py" id="torch.export.dynamic_shapes.ShapesCollection.dynamic_shapes">
<span class="sig-name descname"><span class="pre">dynamic_shapes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">m</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">args</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">kwargs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/dynamic_shapes.html#ShapesCollection.dynamic_shapes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/dynamic_shapes.py#L641"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.dynamic_shapes.ShapesCollection.dynamic_shapes" title="Permalink to this definition">¶</a></dt>
<dd><p>Generate dynamic_shapes.</p>
</dd></dl>

</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes">
<span class="sig-prename descclassname"><span class="pre">torch.export.dynamic_shapes.</span></span><span class="sig-name descname"><span class="pre">refine_dynamic_shapes_from_suggested_fixes</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">msg</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dynamic_shapes</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/dynamic_shapes.html#refine_dynamic_shapes_from_suggested_fixes"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/dynamic_shapes.py#L993"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes" title="Permalink to this definition">¶</a></dt>
<dd><p>For working with export’s dynamic shapes suggested fixes, and/or automatic dynamic shapes.
Refines the given dynamic shapes spec, given a ConstraintViolation error message and the original dynamic shapes.</p>
<p>For most cases behavior is straightforward - i.e. for suggested fixes that specialize or refine a Dim’s range,
or fixes that suggest a derived relation, the new dynamic shapes spec will be updated as such.</p>
<p>e.g.
Suggested fixes:</p>
<blockquote>
<div><p>dim = Dim(‘dim’, min=3, max=6) -&gt; this just refines the dim’s range
dim = 4 -&gt; this specializes to a constant
dy = dx + 1 -&gt; dy was specified as an independent dim, but is actually tied to dx with this relation</p>
</div></blockquote>
<p>However, suggested fixes associated with derived dims can be more complicated.
For example, if a suggested fix is provided for a root dim, the new derived dim value is evaluated based on the root.</p>
<p>e.g.
dx = Dim(‘dx’)
dy = dx + 2
dynamic_shapes = {“x”: (dx,), “y”: (dy,)}</p>
<p>Suggested fixes:</p>
<blockquote>
<div><p>dx = 4  # specialization will lead to dy also specializing = 6
dx = Dim(‘dx’, max=6)  # dy now has max = 8</p>
</div></blockquote>
<p>Derived dims suggested fixes can also be used to express divisibility constraints.
This involves creating new root dims that aren’t tied to a particular input shape.
In this case the root dims won’t appear directly in the new spec, but as a root of
one of the dims.</p>
<p>e.g.
Suggested fixes:</p>
<blockquote>
<div><p>_dx = Dim(‘_dx’, max=1024)  # this won’t appear in the return result, but dx will
dx = 4*_dx  # dx is now divisible by 4, with a max value of 4096</p>
</div></blockquote>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Union" title="(in Python v3.13)"><em>Union</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.13)"><em>Dict</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>, <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a>], <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.13)"><em>Tuple</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a>], <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.List" title="(in Python v3.13)"><em>List</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a>]]</p>
</dd>
</dl>
</dd></dl>

<dl class="py attribute">
<dt class="sig sig-object py" id="torch.export.Constraint">
<span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">Constraint</span></span><a class="headerlink" href="#torch.export.Constraint" title="Permalink to this definition">¶</a></dt>
<dd><p>alias of <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Union" title="(in Python v3.13)"><code class="xref py py-obj docutils literal notranslate"><span class="pre">Union</span></code></a>[<code class="xref py py-class docutils literal notranslate"><span class="pre">_Constraint</span></code>, <code class="xref py py-class docutils literal notranslate"><span class="pre">_DerivedConstraint</span></code>, <code class="xref py py-class docutils literal notranslate"><span class="pre">_RelaxedConstraint</span></code>]</p>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.ExportedProgram">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">ExportedProgram</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">root</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">graph</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">graph_signature</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">state_dict</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">range_constraints</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">module_call_graph</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">example_inputs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">constants</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">verifiers</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L781"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram" title="Permalink to this definition">¶</a></dt>
<dd><p>Package of a program from <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a>. It contains
an <a class="reference internal" href="fx.html#torch.fx.Graph" title="torch.fx.Graph"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.fx.Graph</span></code></a> that represents Tensor computation, a state_dict containing
tensor values of all lifted parameters and buffers, and various metadata.</p>
<p>You can call an ExportedProgram like the original callable traced by
<a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a> with the same calling convention.</p>
<p>To perform transformations on the graph, use <code class="docutils literal notranslate"><span class="pre">.module</span></code> property to access
an <a class="reference internal" href="fx.html#torch.fx.GraphModule" title="torch.fx.GraphModule"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.fx.GraphModule</span></code></a>. You can then use
<a class="reference external" href="https://pytorch.org/docs/stable/fx.html#writing-transformations">FX transformation</a>
to rewrite the graph. Afterwards, you can simply use <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a>
again to construct a correct ExportedProgram.</p>
<dl class="field-list simple">
</dl>
<dl class="py method">
<dt class="sig sig-object py" id="torch.export.ExportedProgram.module">
<span class="sig-name descname"><span class="pre">module</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram.module"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L1178"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram.module" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns a self contained GraphModule with all the parameters/buffers inlined.</p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><p id="torch.nn.Module"/><a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.modules.module.Module"><em>Module</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.ExportedProgram.buffers">
<span class="sig-name descname"><span class="pre">buffers</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram.buffers"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L893"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram.buffers" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns an iterator over original module buffers.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This API is experimental and is <em>NOT</em> backward-compatible.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Iterator" title="(in Python v3.13)"><em>Iterator</em></a>[<a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><em>Tensor</em></a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.ExportedProgram.named_buffers">
<span class="sig-name descname"><span class="pre">named_buffers</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram.named_buffers"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L901"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram.named_buffers" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns an iterator over original module buffers, yielding
both the name of the buffer as well as the buffer itself.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This API is experimental and is <em>NOT</em> backward-compatible.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Iterator" title="(in Python v3.13)"><em>Iterator</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.13)"><em>Tuple</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>, <a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><em>Tensor</em></a>]]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.ExportedProgram.parameters">
<span class="sig-name descname"><span class="pre">parameters</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram.parameters"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L876"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram.parameters" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns an iterator over original module’s parameters.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This API is experimental and is <em>NOT</em> backward-compatible.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Iterator" title="(in Python v3.13)"><em>Iterator</em></a>[<a class="reference internal" href="generated/torch.nn.parameter.Parameter.html#torch.nn.parameter.Parameter" title="torch.nn.parameter.Parameter"><em>Parameter</em></a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.ExportedProgram.named_parameters">
<span class="sig-name descname"><span class="pre">named_parameters</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram.named_parameters"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L884"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram.named_parameters" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns an iterator over original module parameters, yielding
both the name of the parameter as well as the parameter itself.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>This API is experimental and is <em>NOT</em> backward-compatible.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Iterator" title="(in Python v3.13)"><em>Iterator</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Tuple" title="(in Python v3.13)"><em>Tuple</em></a>[<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)">str</a>, <a class="reference internal" href="generated/torch.nn.parameter.Parameter.html#torch.nn.parameter.Parameter" title="torch.nn.parameter.Parameter"><em>Parameter</em></a>]]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.ExportedProgram.run_decompositions">
<span class="sig-name descname"><span class="pre">run_decompositions</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">decomp_table</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ExportedProgram.run_decompositions"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L1206"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportedProgram.run_decompositions" title="Permalink to this definition">¶</a></dt>
<dd><p>Run a set of decompositions on the exported program and returns a new
exported program. By default we will run the Core ATen decompositions to
get operators in the
<a class="reference external" href="https://pytorch.org/docs/stable/torch.compiler_ir.html">Core ATen Operator Set</a>.</p>
<p>For now, we do not decompose joint graphs.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><p><strong>decomp_table</strong> (<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Optional" title="(in Python v3.13)"><em>Optional</em></a><em>[</em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.13)"><em>Dict</em></a><em>[</em><em>OperatorBase</em><em>, </em><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Callable" title="(in Python v3.13)"><em>Callable</em></a><em>]</em><em>]</em>) – An optional argument that specifies decomp behaviour for Aten ops
(1) If None, we decompose to core aten decompositions
(2) If empty, we don’t decompose any operator</p>
</dd>
<dt class="field-even">Return type</dt>
<dd class="field-even"><p><a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.exported_program.ExportedProgram"><em>ExportedProgram</em></a></p>
</dd>
</dl>
<p>Some examples:</p>
<p>If you don’t want to decompose anything</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="o">...</span><span class="p">)</span>
<span class="n">ep</span> <span class="o">=</span> <span class="n">ep</span><span class="o">.</span><span class="n">run_decompositions</span><span class="p">(</span><span class="n">decomp_table</span><span class="o">=</span><span class="p">{})</span>
</pre></div>
</div>
<p>If you want to get a core aten operator set except for certain operator, you can do following:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">ep</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">export</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="o">...</span><span class="p">)</span>
<span class="n">decomp_table</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">export</span><span class="o">.</span><span class="n">default_decompositions</span><span class="p">()</span>
<span class="n">decomp_table</span><span class="p">[</span><span class="n">your_op</span><span class="p">]</span> <span class="o">=</span> <span class="n">your_custom_decomp</span>
<span class="n">ep</span> <span class="o">=</span> <span class="n">ep</span><span class="o">.</span><span class="n">run_decompositions</span><span class="p">(</span><span class="n">decomp_table</span><span class="o">=</span><span class="n">decomp_table</span><span class="p">)</span>
</pre></div>
</div>
</dd></dl>

</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.ExportBackwardSignature">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">ExportBackwardSignature</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">gradients_to_parameters</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">gradients_to_user_inputs</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">loss_output</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#ExportBackwardSignature"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L144"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportBackwardSignature" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.ExportGraphSignature">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">ExportGraphSignature</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">input_specs</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">output_specs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#ExportGraphSignature"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L151"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ExportGraphSignature" title="Permalink to this definition">¶</a></dt>
<dd><p><a class="reference internal" href="#torch.export.ExportGraphSignature" title="torch.export.ExportGraphSignature"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportGraphSignature</span></code></a> models the input/output signature of Export Graph,
which is a fx.Graph with stronger invariants gurantees.</p>
<p>Export Graph is functional and does not access “states” like parameters
or buffers within the graph via <code class="docutils literal notranslate"><span class="pre">getattr</span></code> nodes. Instead, <a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code></a>
gurantees that parameters, buffers, and constant tensors are lifted out of
the graph as inputs.  Similarly, any mutations to buffers are not included
in the graph either, instead the updated values of mutated buffers are
modeled as additional outputs of Export Graph.</p>
<p>The ordering of all inputs and outputs are:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">Inputs</span> <span class="o">=</span> <span class="p">[</span><span class="o">*</span><span class="n">parameters_buffers_constant_tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">flattened_user_inputs</span><span class="p">]</span>
<span class="n">Outputs</span> <span class="o">=</span> <span class="p">[</span><span class="o">*</span><span class="n">mutated_inputs</span><span class="p">,</span> <span class="o">*</span><span class="n">flattened_user_outputs</span><span class="p">]</span>
</pre></div>
</div>
<p>e.g. If following module is exported:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">CustomModule</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</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="nb">super</span><span class="p">(</span><span class="n">CustomModule</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="c1"># Define a parameter</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">my_parameter</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</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="mf">2.0</span><span class="p">))</span>

        <span class="c1"># Define two buffers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">register_buffer</span><span class="p">(</span><span class="s1">&#39;my_buffer1&#39;</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="mf">3.0</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">register_buffer</span><span class="p">(</span><span class="s1">&#39;my_buffer2&#39;</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="mf">4.0</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">):</span>
        <span class="c1"># Use the parameter, buffers, and both inputs in the forward method</span>
        <span class="n">output</span> <span class="o">=</span> <span class="p">(</span><span class="n">x1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">my_parameter</span><span class="p">)</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">my_buffer1</span> <span class="o">+</span> <span class="n">x2</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">my_buffer2</span>

        <span class="c1"># Mutate one of the buffers (e.g., increment it by 1)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">my_buffer2</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="mf">1.0</span><span class="p">)</span> <span class="c1"># In-place addition</span>

        <span class="k">return</span> <span class="n">output</span>
</pre></div>
</div>
<p>Resulting Graph would be:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">graph</span><span class="p">():</span>
    <span class="o">%</span><span class="n">arg0_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg0_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg1_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg1_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg2_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg2_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg3_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg3_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg4_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg4_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">add_tensor</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">arg3_1</span><span class="p">,</span> <span class="o">%</span><span class="n">arg0_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">mul_tensor</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">add_tensor</span><span class="p">,</span> <span class="o">%</span><span class="n">arg1_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">mul_tensor_1</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">arg4_1</span><span class="p">,</span> <span class="o">%</span><span class="n">arg2_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">add_tensor_1</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">mul_tensor</span><span class="p">,</span> <span class="o">%</span><span class="n">mul_tensor_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">add_tensor_2</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">arg2_1</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">add_tensor_2</span><span class="p">,</span> <span class="n">add_tensor_1</span><span class="p">)</span>
</pre></div>
</div>
<p>Resulting ExportGraphSignature would be:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportGraphSignature</span><span class="p">(</span>
    <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg0_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_parameter&#39;</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg1_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_buffer1&#39;</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg2_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_buffer2&#39;</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg3_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg4_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>
    <span class="p">],</span>
    <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
        <span class="n">OutputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add_2&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_buffer2&#39;</span><span class="p">),</span>
        <span class="n">OutputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>
    <span class="p">]</span>
<span class="p">)</span>
</pre></div>
</div>
<dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.ModuleCallSignature">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">ModuleCallSignature</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">inputs</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">torch.export.graph_signature.TensorArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymIntArgument" title="torch.export.graph_signature.SymIntArgument"><span class="pre">torch.export.graph_signature.SymIntArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymFloatArgument" title="torch.export.graph_signature.SymFloatArgument"><span class="pre">torch.export.graph_signature.SymFloatArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymBoolArgument" title="torch.export.graph_signature.SymBoolArgument"><span class="pre">torch.export.graph_signature.SymBoolArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.ConstantArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.CustomObjArgument" title="torch.export.graph_signature.CustomObjArgument"><span class="pre">torch.export.graph_signature.CustomObjArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.TokenArgument</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">outputs</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">torch.export.graph_signature.TensorArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymIntArgument" title="torch.export.graph_signature.SymIntArgument"><span class="pre">torch.export.graph_signature.SymIntArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymFloatArgument" title="torch.export.graph_signature.SymFloatArgument"><span class="pre">torch.export.graph_signature.SymFloatArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymBoolArgument" title="torch.export.graph_signature.SymBoolArgument"><span class="pre">torch.export.graph_signature.SymBoolArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.ConstantArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.CustomObjArgument" title="torch.export.graph_signature.CustomObjArgument"><span class="pre">torch.export.graph_signature.CustomObjArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.TokenArgument</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">in_spec</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">torch.utils._pytree.TreeSpec</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">out_spec</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">torch.utils._pytree.TreeSpec</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">forward_arg_names</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">List</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ModuleCallSignature"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L101"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ModuleCallSignature" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.ModuleCallEntry">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.</span></span><span class="sig-name descname"><span class="pre">ModuleCallEntry</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fqn</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em>, <em class="sig-param"><span class="n"><span class="pre">signature</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><a class="reference internal" href="#torch.export.ModuleCallSignature" title="torch.export.exported_program.ModuleCallSignature"><span class="pre">torch.export.exported_program.ModuleCallSignature</span></a><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/exported_program.html#ModuleCallEntry"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/exported_program.py#L118"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.ModuleCallEntry" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<span class="target" id="module-torch.export.decomp_utils"></span><dl class="py class">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.decomp_utils.</span></span><span class="sig-name descname"><span class="pre">CustomDecompTable</span></span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L16"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable" title="Permalink to this definition">¶</a></dt>
<dd><p>This is a custom dictionary that is specifically used for handling decomp_table in export.
The reason we need this is because in the new world, you can only <em>delete</em> an op from decomp
table to preserve it. This is problematic for custom ops because we don’t know when the custom
op will actually be loaded to the dispatcher. As a result, we need to record the custom ops operations
until we really need to materialize it (which is when we run decomposition pass.)</p>
<dl class="simple">
<dt>Invariants we hold are:</dt><dd><ol class="arabic simple">
<li><p>All aten decomp is loaded at the init time</p></li>
<li><p>We materialize ALL ops when user ever reads from the table to make it more likely
that dispatcher picks up the custom op.</p></li>
<li><p>If it is write operation, we don’t necessarily materialize</p></li>
<li><p>We load the final time during export, right before calling run_decompositions()</p></li>
</ol>
</dd>
</dl>
<dl class="py method">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable.copy">
<span class="sig-name descname"><span class="pre">copy</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable.copy"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L88"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable.copy" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable" title="torch.export.decomp_utils.CustomDecompTable"><em>CustomDecompTable</em></a></p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable.items">
<span class="sig-name descname"><span class="pre">items</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable.items"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L125"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable.items" title="Permalink to this definition">¶</a></dt>
<dd></dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable.keys">
<span class="sig-name descname"><span class="pre">keys</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable.keys"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L58"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable.keys" title="Permalink to this definition">¶</a></dt>
<dd></dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable.materialize">
<span class="sig-name descname"><span class="pre">materialize</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable.materialize"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L129"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable.materialize" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Dict" title="(in Python v3.13)"><em>Dict</em></a>[<em>OperatorBase</em>, <a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Callable" title="(in Python v3.13)"><em>Callable</em></a>]</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable.pop">
<span class="sig-name descname"><span class="pre">pop</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="o"><span class="pre">*</span></span><span class="n"><span class="pre">args</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable.pop"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L95"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable.pop" title="Permalink to this definition">¶</a></dt>
<dd></dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.decomp_utils.CustomDecompTable.update">
<span class="sig-name descname"><span class="pre">update</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">other_dict</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/decomp_utils.html#CustomDecompTable.update"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/decomp_utils.py#L65"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.decomp_utils.CustomDecompTable.update" title="Permalink to this definition">¶</a></dt>
<dd></dd></dl>

</dd></dl>

<span class="target" id="module-torch.export.exported_program"></span><span class="target" id="module-torch.export.graph_signature"></span><dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.InputKind">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">InputKind</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">value</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#InputKind"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L78"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.InputKind" title="Permalink to this definition">¶</a></dt>
<dd><p>An enumeration.</p>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.InputSpec">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">InputSpec</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">kind</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference internal" href="#torch.export.graph_signature.InputKind" title="torch.export.graph_signature.InputKind"><span class="pre">torch.export.graph_signature.InputKind</span></a></span></em>, <em class="sig-param"><span class="n"><span class="pre">arg</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">torch.export.graph_signature.TensorArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymIntArgument" title="torch.export.graph_signature.SymIntArgument"><span class="pre">torch.export.graph_signature.SymIntArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymFloatArgument" title="torch.export.graph_signature.SymFloatArgument"><span class="pre">torch.export.graph_signature.SymFloatArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymBoolArgument" title="torch.export.graph_signature.SymBoolArgument"><span class="pre">torch.export.graph_signature.SymBoolArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.ConstantArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.CustomObjArgument" title="torch.export.graph_signature.CustomObjArgument"><span class="pre">torch.export.graph_signature.CustomObjArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.TokenArgument</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">persistent</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.13)"><span class="pre">bool</span></a><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#InputSpec"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L87"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.InputSpec" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.OutputKind">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">OutputKind</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">value</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#OutputKind"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L113"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.OutputKind" title="Permalink to this definition">¶</a></dt>
<dd><p>An enumeration.</p>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.OutputSpec">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">OutputSpec</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">kind</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference internal" href="#torch.export.graph_signature.OutputKind" title="torch.export.graph_signature.OutputKind"><span class="pre">torch.export.graph_signature.OutputKind</span></a></span></em>, <em class="sig-param"><span class="n"><span class="pre">arg</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">torch.export.graph_signature.TensorArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymIntArgument" title="torch.export.graph_signature.SymIntArgument"><span class="pre">torch.export.graph_signature.SymIntArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymFloatArgument" title="torch.export.graph_signature.SymFloatArgument"><span class="pre">torch.export.graph_signature.SymFloatArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.SymBoolArgument" title="torch.export.graph_signature.SymBoolArgument"><span class="pre">torch.export.graph_signature.SymBoolArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.ConstantArgument</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><a class="reference internal" href="#torch.export.graph_signature.CustomObjArgument" title="torch.export.graph_signature.CustomObjArgument"><span class="pre">torch.export.graph_signature.CustomObjArgument</span></a><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">torch.export.graph_signature.TokenArgument</span><span class="p"><span class="pre">]</span></span></span></em>, <em class="sig-param"><span class="n"><span class="pre">target</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a><span class="p"><span class="pre">]</span></span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#OutputSpec"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L123"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.OutputSpec" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.SymIntArgument">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">SymIntArgument</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">name</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#SymIntArgument"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L39"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.SymIntArgument" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.SymBoolArgument">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">SymBoolArgument</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">name</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#SymBoolArgument"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L49"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.SymBoolArgument" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.SymFloatArgument">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">SymFloatArgument</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">name</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#SymFloatArgument"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L44"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.SymFloatArgument" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.ExportGraphSignature">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">ExportGraphSignature</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">input_specs</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">output_specs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#ExportGraphSignature"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L151"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.ExportGraphSignature" title="Permalink to this definition">¶</a></dt>
<dd><p><a class="reference internal" href="#torch.export.graph_signature.ExportGraphSignature" title="torch.export.graph_signature.ExportGraphSignature"><code class="xref py py-class docutils literal notranslate"><span class="pre">ExportGraphSignature</span></code></a> models the input/output signature of Export Graph,
which is a fx.Graph with stronger invariants gurantees.</p>
<p>Export Graph is functional and does not access “states” like parameters
or buffers within the graph via <code class="docutils literal notranslate"><span class="pre">getattr</span></code> nodes. Instead, <code class="xref py py-func docutils literal notranslate"><span class="pre">export()</span></code>
gurantees that parameters, buffers, and constant tensors are lifted out of
the graph as inputs.  Similarly, any mutations to buffers are not included
in the graph either, instead the updated values of mutated buffers are
modeled as additional outputs of Export Graph.</p>
<p>The ordering of all inputs and outputs are:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">Inputs</span> <span class="o">=</span> <span class="p">[</span><span class="o">*</span><span class="n">parameters_buffers_constant_tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">flattened_user_inputs</span><span class="p">]</span>
<span class="n">Outputs</span> <span class="o">=</span> <span class="p">[</span><span class="o">*</span><span class="n">mutated_inputs</span><span class="p">,</span> <span class="o">*</span><span class="n">flattened_user_outputs</span><span class="p">]</span>
</pre></div>
</div>
<p>e.g. If following module is exported:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">CustomModule</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</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="nb">super</span><span class="p">(</span><span class="n">CustomModule</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="c1"># Define a parameter</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">my_parameter</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Parameter</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="mf">2.0</span><span class="p">))</span>

        <span class="c1"># Define two buffers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">register_buffer</span><span class="p">(</span><span class="s1">&#39;my_buffer1&#39;</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="mf">3.0</span><span class="p">))</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">register_buffer</span><span class="p">(</span><span class="s1">&#39;my_buffer2&#39;</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="mf">4.0</span><span class="p">))</span>

    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">):</span>
        <span class="c1"># Use the parameter, buffers, and both inputs in the forward method</span>
        <span class="n">output</span> <span class="o">=</span> <span class="p">(</span><span class="n">x1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">my_parameter</span><span class="p">)</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">my_buffer1</span> <span class="o">+</span> <span class="n">x2</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">my_buffer2</span>

        <span class="c1"># Mutate one of the buffers (e.g., increment it by 1)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">my_buffer2</span><span class="o">.</span><span class="n">add_</span><span class="p">(</span><span class="mf">1.0</span><span class="p">)</span> <span class="c1"># In-place addition</span>

        <span class="k">return</span> <span class="n">output</span>
</pre></div>
</div>
<p>Resulting Graph would be:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">graph</span><span class="p">():</span>
    <span class="o">%</span><span class="n">arg0_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg0_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg1_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg1_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg2_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg2_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg3_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg3_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">arg4_1</span> <span class="o">:=</span> <span class="n">placeholder</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">arg4_1</span><span class="p">]</span>
    <span class="o">%</span><span class="n">add_tensor</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">arg3_1</span><span class="p">,</span> <span class="o">%</span><span class="n">arg0_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">mul_tensor</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">add_tensor</span><span class="p">,</span> <span class="o">%</span><span class="n">arg1_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">mul_tensor_1</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">mul</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">arg4_1</span><span class="p">,</span> <span class="o">%</span><span class="n">arg2_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">add_tensor_1</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">mul_tensor</span><span class="p">,</span> <span class="o">%</span><span class="n">mul_tensor_1</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="o">%</span><span class="n">add_tensor_2</span> <span class="o">:=</span> <span class="n">call_function</span><span class="p">[</span><span class="n">target</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">ops</span><span class="o">.</span><span class="n">aten</span><span class="o">.</span><span class="n">add</span><span class="o">.</span><span class="n">Tensor</span><span class="p">](</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="o">%</span><span class="n">arg2_1</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">),</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{})</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">add_tensor_2</span><span class="p">,</span> <span class="n">add_tensor_1</span><span class="p">)</span>
</pre></div>
</div>
<p>Resulting ExportGraphSignature would be:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">ExportGraphSignature</span><span class="p">(</span>
    <span class="n">input_specs</span><span class="o">=</span><span class="p">[</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">PARAMETER</span><span class="p">:</span> <span class="mi">2</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg0_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_parameter&#39;</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg1_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_buffer1&#39;</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">BUFFER</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg2_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_buffer2&#39;</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg3_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">),</span>
        <span class="n">InputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">InputKind</span><span class="o">.</span><span class="n">USER_INPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;arg4_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>
    <span class="p">],</span>
    <span class="n">output_specs</span><span class="o">=</span><span class="p">[</span>
        <span class="n">OutputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">BUFFER_MUTATION</span><span class="p">:</span> <span class="mi">3</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add_2&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="s1">&#39;my_buffer2&#39;</span><span class="p">),</span>
        <span class="n">OutputSpec</span><span class="p">(</span><span class="n">kind</span><span class="o">=&lt;</span><span class="n">OutputKind</span><span class="o">.</span><span class="n">USER_OUTPUT</span><span class="p">:</span> <span class="mi">1</span><span class="o">&gt;</span><span class="p">,</span> <span class="n">arg</span><span class="o">=</span><span class="n">TensorArgument</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">&#39;add_1&#39;</span><span class="p">),</span> <span class="n">target</span><span class="o">=</span><span class="kc">None</span><span class="p">)</span>
    <span class="p">]</span>
<span class="p">)</span>
</pre></div>
</div>
<dl class="field-list simple">
</dl>
<dl class="py method">
<dt class="sig sig-object py" id="torch.export.graph_signature.ExportGraphSignature.replace_all_uses">
<span class="sig-name descname"><span class="pre">replace_all_uses</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">old</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">new</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#ExportGraphSignature.replace_all_uses"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L456"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.ExportGraphSignature.replace_all_uses" title="Permalink to this definition">¶</a></dt>
<dd><p>Replace all uses of the old name with new name in the signature.</p>
<dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.export.graph_signature.ExportGraphSignature.get_replace_hook">
<span class="sig-name descname"><span class="pre">get_replace_hook</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">replace_inputs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#ExportGraphSignature.get_replace_hook"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L479"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.ExportGraphSignature.get_replace_hook" title="Permalink to this definition">¶</a></dt>
<dd></dd></dl>

</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.graph_signature.CustomObjArgument">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.graph_signature.</span></span><span class="sig-name descname"><span class="pre">CustomObjArgument</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">name</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em>, <em class="sig-param"><span class="n"><span class="pre">class_fqn</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><span class="pre">str</span></a></span></em>, <em class="sig-param"><span class="n"><span class="pre">fake_val</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">torch._library.fake_class_registry.FakeScriptObject</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/graph_signature.html#CustomObjArgument"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/graph_signature.py#L54"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.graph_signature.CustomObjArgument" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
</dl>
</dd></dl>

<span class="target" id="module-torch.export.dynamic_shapes"></span><span class="target" id="module-torch.export.unflatten"></span><dl class="py class">
<dt class="sig sig-object py" id="torch.export.unflatten.FlatArgsAdapter">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.unflatten.</span></span><span class="sig-name descname"><span class="pre">FlatArgsAdapter</span></span><a class="reference internal" href="_modules/torch/export/unflatten.html#FlatArgsAdapter"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/unflatten.py#L256"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.unflatten.FlatArgsAdapter" title="Permalink to this definition">¶</a></dt>
<dd><p>Adapts input arguments with <code class="docutils literal notranslate"><span class="pre">input_spec</span></code> to align <code class="docutils literal notranslate"><span class="pre">target_spec</span></code>.</p>
<dl class="py method">
<dt class="sig sig-object py" id="torch.export.unflatten.FlatArgsAdapter.adapt">
<em class="property"><span class="pre">abstract</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">adapt</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">target_spec</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_spec</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">input_args</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/unflatten.html#FlatArgsAdapter.adapt"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/unflatten.py#L261"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.unflatten.FlatArgsAdapter.adapt" title="Permalink to this definition">¶</a></dt>
<dd><p>NOTE: This adapter may mutate given <code class="docutils literal notranslate"><span class="pre">input_args_with_path</span></code>.</p>
<dl class="field-list simple">
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.List" title="(in Python v3.13)"><em>List</em></a>[<a class="reference external" href="https://docs.python.org/3/library/typing.html#typing.Any" title="(in Python v3.13)"><em>Any</em></a>]</p>
</dd>
</dl>
</dd></dl>

</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.unflatten.InterpreterModule">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.unflatten.</span></span><span class="sig-name descname"><span class="pre">InterpreterModule</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">graph</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/unflatten.html#InterpreterModule"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/unflatten.py#L123"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.unflatten.InterpreterModule" title="Permalink to this definition">¶</a></dt>
<dd><p>A module that uses torch.fx.Interpreter to execute instead of the usual
codegen that GraphModule uses. This provides better stack trace information
and makes it easier to debug execution.</p>
<dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.export.unflatten.InterpreterModuleDispatcher">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.export.unflatten.</span></span><span class="sig-name descname"><span class="pre">InterpreterModuleDispatcher</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">attrs</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">call_modules</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/unflatten.html#InterpreterModuleDispatcher"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/unflatten.py#L209"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.unflatten.InterpreterModuleDispatcher" title="Permalink to this definition">¶</a></dt>
<dd><p>A module that carries a sequence of InterpreterModules corresponding to
a sequence of calls of that module. Each call to the module dispatches
to the next InterpreterModule, and wraps back around after the last.</p>
<dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.export.unflatten.unflatten">
<span class="sig-prename descclassname"><span class="pre">torch.export.unflatten.</span></span><span class="sig-name descname"><span class="pre">unflatten</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">module</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">flat_args_adapter</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/unflatten.html#unflatten"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/unflatten.py#L672"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.unflatten.unflatten" title="Permalink to this definition">¶</a></dt>
<dd><p>Unflatten an ExportedProgram, producing a module with the same module
hierarchy as the original eager module. This can be useful if you are trying
to use <a class="reference internal" href="#module-torch.export" title="torch.export"><code class="xref py py-mod docutils literal notranslate"><span class="pre">torch.export</span></code></a> with another system that expects a module
hierachy instead of the flat graph that <a class="reference internal" href="#module-torch.export" title="torch.export"><code class="xref py py-mod docutils literal notranslate"><span class="pre">torch.export</span></code></a> usually produces.</p>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>The args/kwargs of unflattened modules will not necessarily match
the eager module, so doing a module swap (e.g. <code class="code docutils literal notranslate"><span class="pre">self.submod</span> <span class="pre">=</span>
<span class="pre">new_mod</span></code>) will not necessarily work. If you need to swap a module out, you
need to set the <code class="code docutils literal notranslate"><span class="pre">preserve_module_call_signature</span></code> parameter of
<a class="reference internal" href="#torch.export.export" title="torch.export.export"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.export.export()</span></code></a>.</p>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>module</strong> (<a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><em>ExportedProgram</em></a>) – The ExportedProgram to unflatten.</p></li>
<li><p><strong>flat_args_adapter</strong> (<em>Optional</em><em>[</em><a class="reference internal" href="#torch.export.unflatten.FlatArgsAdapter" title="torch.export.unflatten.FlatArgsAdapter"><em>FlatArgsAdapter</em></a><em>]</em>) – Adapt flat args if input TreeSpec does not match with exported module’s.</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>An instance of <code class="xref py py-class docutils literal notranslate"><span class="pre">UnflattenedModule</span></code>, which has the same module
hierarchy as the original eager module pre-export.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><em>UnflattenedModule</em></p>
</dd>
</dl>
</dd></dl>

<span class="target" id="module-torch.export.custom_obj"></span><span class="target" id="module-torch.export.experimental"></span><span class="target" id="module-torch.export.passes"></span><dl class="py function">
<dt class="sig sig-object py" id="torch.export.passes.move_to_device_pass">
<span class="sig-prename descclassname"><span class="pre">torch.export.passes.</span></span><span class="sig-name descname"><span class="pre">move_to_device_pass</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">ep</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">location</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/export/passes.html#move_to_device_pass"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="reference external" href="https://github.com/pytorch/pytorch/blob/v2.6.0/torch/export/passes/__init__.py#L11"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.export.passes.move_to_device_pass" title="Permalink to this definition">¶</a></dt>
<dd><p>Move the exported program to the given device.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters</dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>ep</strong> (<a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram"><em>ExportedProgram</em></a>) – The exported program to move.</p></li>
<li><p><strong>location</strong> (<em>Union</em><em>[</em><a class="reference internal" href="tensor_attributes.html#torch.device" title="torch.device"><em>torch.device</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><em>Dict</em><em>[</em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>, </em><a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.13)"><em>str</em></a><em>]</em><em>]</em>) – The device to move the exported program to.
If a string, it is interpreted as a device name.
If a dict, it is interpreted as a mapping from
the existing device to the intended one</p></li>
</ul>
</dd>
<dt class="field-even">Returns</dt>
<dd class="field-even"><p>The moved exported program.</p>
</dd>
<dt class="field-odd">Return type</dt>
<dd class="field-odd"><p><a class="reference internal" href="#torch.export.ExportedProgram" title="torch.export.ExportedProgram">ExportedProgram</a></p>
</dd>
</dl>
</dd></dl>

</div>
</div>


             </article>
             
            </div>
            <footer>
  
    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="export.ir_spec.html" class="btn btn-neutral float-right" title="torch.export IR Specification" accesskey="n" rel="next">Next <img src="_static/images/chevron-right-orange.svg" class="next-page"></a>
      
      
        <a href="backends.html" class="btn btn-neutral" title="torch.backends" accesskey="p" rel="prev"><img src="_static/images/chevron-right-orange.svg" class="previous-page"> Previous</a>
      
    </div>
  

    <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">
              <ul>
<li><a class="reference internal" href="#">torch.export</a><ul>
<li><a class="reference internal" href="#overview">Overview</a><ul>
<li><a class="reference internal" href="#existing-frameworks">Existing frameworks</a></li>
</ul>
</li>
<li><a class="reference internal" href="#exporting-a-pytorch-model">Exporting a PyTorch Model</a><ul>
<li><a class="reference internal" href="#an-example">An Example</a></li>
<li><a class="reference internal" href="#non-strict-export">Non-Strict Export</a></li>
<li><a class="reference internal" href="#export-for-training-and-inference">Export for Training and Inference</a></li>
<li><a class="reference internal" href="#expressing-dynamism">Expressing Dynamism</a></li>
<li><a class="reference internal" href="#serialization">Serialization</a></li>
<li><a class="reference internal" href="#specializations">Specializations</a><ul>
<li><a class="reference internal" href="#input-tensor-shapes">Input Tensor Shapes</a></li>
<li><a class="reference internal" href="#python-primitives">Python Primitives</a></li>
<li><a class="reference internal" href="#python-containers">Python Containers</a></li>
</ul>
</li>
</ul>
</li>
<li><a class="reference internal" href="#limitations-of-torch-export">Limitations of torch.export</a><ul>
<li><a class="reference internal" href="#graph-breaks">Graph Breaks</a></li>
<li><a class="reference internal" href="#data-shape-dependent-control-flow">Data/Shape-Dependent Control Flow</a></li>
<li><a class="reference internal" href="#missing-fake-meta-abstract-kernels-for-operators">Missing Fake/Meta/Abstract Kernels for Operators</a></li>
</ul>
</li>
<li><a class="reference internal" href="#read-more">Read More</a><ul>
</ul>
</li>
<li><a class="reference internal" href="#module-torch.export">API Reference</a><ul>
<li><a class="reference internal" href="#torch.export.export"><code class="docutils literal notranslate"><span class="pre">export()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.save"><code class="docutils literal notranslate"><span class="pre">save()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.load"><code class="docutils literal notranslate"><span class="pre">load()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.register_dataclass"><code class="docutils literal notranslate"><span class="pre">register_dataclass()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.dynamic_shapes.Dim"><code class="docutils literal notranslate"><span class="pre">Dim()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.exported_program.default_decompositions"><code class="docutils literal notranslate"><span class="pre">default_decompositions()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.dims"><code class="docutils literal notranslate"><span class="pre">dims()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.dynamic_shapes.ShapesCollection"><code class="docutils literal notranslate"><span class="pre">ShapesCollection</span></code></a><ul>
<li><a class="reference internal" href="#torch.export.dynamic_shapes.ShapesCollection.dynamic_shapes"><code class="docutils literal notranslate"><span class="pre">ShapesCollection.dynamic_shapes()</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes"><code class="docutils literal notranslate"><span class="pre">refine_dynamic_shapes_from_suggested_fixes()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.Constraint"><code class="docutils literal notranslate"><span class="pre">Constraint</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportedProgram"><code class="docutils literal notranslate"><span class="pre">ExportedProgram</span></code></a><ul>
<li><a class="reference internal" href="#torch.export.ExportedProgram.module"><code class="docutils literal notranslate"><span class="pre">ExportedProgram.module()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportedProgram.buffers"><code class="docutils literal notranslate"><span class="pre">ExportedProgram.buffers()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportedProgram.named_buffers"><code class="docutils literal notranslate"><span class="pre">ExportedProgram.named_buffers()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportedProgram.parameters"><code class="docutils literal notranslate"><span class="pre">ExportedProgram.parameters()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportedProgram.named_parameters"><code class="docutils literal notranslate"><span class="pre">ExportedProgram.named_parameters()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportedProgram.run_decompositions"><code class="docutils literal notranslate"><span class="pre">ExportedProgram.run_decompositions()</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#torch.export.ExportBackwardSignature"><code class="docutils literal notranslate"><span class="pre">ExportBackwardSignature</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ExportGraphSignature"><code class="docutils literal notranslate"><span class="pre">ExportGraphSignature</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ModuleCallSignature"><code class="docutils literal notranslate"><span class="pre">ModuleCallSignature</span></code></a></li>
<li><a class="reference internal" href="#torch.export.ModuleCallEntry"><code class="docutils literal notranslate"><span class="pre">ModuleCallEntry</span></code></a></li>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable</span></code></a><ul>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable.copy"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable.copy()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable.items"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable.items()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable.keys"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable.keys()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable.materialize"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable.materialize()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable.pop"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable.pop()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.decomp_utils.CustomDecompTable.update"><code class="docutils literal notranslate"><span class="pre">CustomDecompTable.update()</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#torch.export.graph_signature.InputKind"><code class="docutils literal notranslate"><span class="pre">InputKind</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.InputSpec"><code class="docutils literal notranslate"><span class="pre">InputSpec</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.OutputKind"><code class="docutils literal notranslate"><span class="pre">OutputKind</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.OutputSpec"><code class="docutils literal notranslate"><span class="pre">OutputSpec</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.SymIntArgument"><code class="docutils literal notranslate"><span class="pre">SymIntArgument</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.SymBoolArgument"><code class="docutils literal notranslate"><span class="pre">SymBoolArgument</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.SymFloatArgument"><code class="docutils literal notranslate"><span class="pre">SymFloatArgument</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.ExportGraphSignature"><code class="docutils literal notranslate"><span class="pre">ExportGraphSignature</span></code></a><ul>
<li><a class="reference internal" href="#torch.export.graph_signature.ExportGraphSignature.replace_all_uses"><code class="docutils literal notranslate"><span class="pre">ExportGraphSignature.replace_all_uses()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.graph_signature.ExportGraphSignature.get_replace_hook"><code class="docutils literal notranslate"><span class="pre">ExportGraphSignature.get_replace_hook()</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#torch.export.graph_signature.CustomObjArgument"><code class="docutils literal notranslate"><span class="pre">CustomObjArgument</span></code></a></li>
<li><a class="reference internal" href="#torch.export.unflatten.FlatArgsAdapter"><code class="docutils literal notranslate"><span class="pre">FlatArgsAdapter</span></code></a><ul>
<li><a class="reference internal" href="#torch.export.unflatten.FlatArgsAdapter.adapt"><code class="docutils literal notranslate"><span class="pre">FlatArgsAdapter.adapt()</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#torch.export.unflatten.InterpreterModule"><code class="docutils literal notranslate"><span class="pre">InterpreterModule</span></code></a></li>
<li><a class="reference internal" href="#torch.export.unflatten.InterpreterModuleDispatcher"><code class="docutils literal notranslate"><span class="pre">InterpreterModuleDispatcher</span></code></a></li>
<li><a class="reference internal" href="#torch.export.unflatten.unflatten"><code class="docutils literal notranslate"><span class="pre">unflatten()</span></code></a></li>
<li><a class="reference internal" href="#torch.export.passes.move_to_device_pass"><code class="docutils literal notranslate"><span class="pre">move_to_device_pass()</span></code></a></li>
</ul>
</li>
</ul>
</li>
</ul>

            </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>