2.6/_modules/torch/onnx/utils.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.onnx.utils &mdash; PyTorch 2.6 documentation</title>
  

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

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


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

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

  <!-- Preload the theme fonts -->

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

<!-- Preload the katex fonts -->

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

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

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

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

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

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

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

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

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

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

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

<body class="pytorch-body">

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

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

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

          </div>

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

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

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

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

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

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

  
</div>
        </div>

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

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

        
          <!-- Google Tag Manager (noscript) -->
          <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-T8XT4PS"
          height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>
          <!-- End Google Tag Manager (noscript) -->
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <h1>Source code for torch.onnx.utils</h1><div class="highlight"><pre>
<span></span><span class="c1"># mypy: allow-untyped-defs</span>
<span class="sd">&quot;&quot;&quot;Functions to export models into the ONNX IR format.</span>

<span class="sd">These models can be loaded with the ONNX library and then</span>
<span class="sd">converted to models which run on other deep learning frameworks.</span>
<span class="sd">&quot;&quot;&quot;</span>

<span class="kn">from</span> <span class="nn">__future__</span> <span class="kn">import</span> <span class="n">annotations</span>

<span class="kn">import</span> <span class="nn">contextlib</span>
<span class="kn">import</span> <span class="nn">copy</span>
<span class="kn">import</span> <span class="nn">inspect</span>
<span class="kn">import</span> <span class="nn">re</span>
<span class="kn">import</span> <span class="nn">typing</span>
<span class="kn">import</span> <span class="nn">warnings</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">cast</span><span class="p">,</span> <span class="n">Collection</span><span class="p">,</span> <span class="n">Mapping</span><span class="p">,</span> <span class="n">Sequence</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">torch._C._onnx</span> <span class="k">as</span> <span class="nn">_C_onnx</span>
<span class="kn">import</span> <span class="nn">torch.jit._trace</span>
<span class="kn">import</span> <span class="nn">torch.serialization</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">_C</span>
<span class="kn">from</span> <span class="nn">torch.onnx</span> <span class="kn">import</span> <span class="n">_constants</span><span class="p">,</span> <span class="n">_deprecation</span><span class="p">,</span> <span class="n">errors</span><span class="p">,</span> <span class="n">symbolic_helper</span>  <span class="c1"># noqa: F401</span>
<span class="kn">from</span> <span class="nn">torch.onnx._globals</span> <span class="kn">import</span> <span class="n">GLOBALS</span>
<span class="kn">from</span> <span class="nn">torch.onnx._internal</span> <span class="kn">import</span> <span class="n">diagnostics</span><span class="p">,</span> <span class="n">jit_utils</span><span class="p">,</span> <span class="n">onnx_proto_utils</span><span class="p">,</span> <span class="n">registration</span>


<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s2">&quot;is_in_onnx_export&quot;</span><span class="p">,</span>
    <span class="s2">&quot;select_model_mode_for_export&quot;</span><span class="p">,</span>
    <span class="s2">&quot;disable_apex_o2_state_dict_hook&quot;</span><span class="p">,</span>
    <span class="s2">&quot;setup_onnx_logging&quot;</span><span class="p">,</span>
    <span class="s2">&quot;exporter_context&quot;</span><span class="p">,</span>
    <span class="s2">&quot;export&quot;</span><span class="p">,</span>
    <span class="s2">&quot;model_signature&quot;</span><span class="p">,</span>
    <span class="s2">&quot;warn_on_static_input_change&quot;</span><span class="p">,</span>
    <span class="s2">&quot;unpack_quantized_tensor&quot;</span><span class="p">,</span>
    <span class="s2">&quot;unconvertible_ops&quot;</span><span class="p">,</span>
    <span class="s2">&quot;register_custom_op_symbolic&quot;</span><span class="p">,</span>
    <span class="s2">&quot;unregister_custom_op_symbolic&quot;</span><span class="p">,</span>
<span class="p">]</span>


<div class="viewcode-block" id="is_in_onnx_export"><a class="viewcode-back" href="../../../onnx_torchscript.html#torch.onnx.is_in_onnx_export">[docs]</a><span class="k">def</span> <span class="nf">is_in_onnx_export</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Returns whether it is in the middle of ONNX export.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">in_onnx_export</span></div>


<span class="c1"># TODO(justinchuby): Remove dependency to this global variable from constant_fold.cpp</span>
<span class="c1"># Skip check due to cannot import IValue from torch._C</span>
<span class="n">_params_dict</span> <span class="o">=</span> <span class="p">{}</span>  <span class="c1"># type: ignore[var-annotated]</span>


<div class="viewcode-block" id="select_model_mode_for_export"><a class="viewcode-back" href="../../../onnx_torchscript.html#torch.onnx.select_model_mode_for_export">[docs]</a><span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
<span class="k">def</span> <span class="nf">select_model_mode_for_export</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;A context manager to temporarily set the training mode of ``model``</span>
<span class="sd">    to ``mode``, resetting it when we exit the with-block.</span>

<span class="sd">    Args:</span>
<span class="sd">        model: Same type and meaning as ``model`` arg to :func:`export`.</span>
<span class="sd">        mode: Same type and meaning as ``training`` arg to :func:`export`.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">mode</span><span class="p">,</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span>
            <span class="sa">f</span><span class="s2">&quot;&#39;mode&#39; should be a torch.onnx.TrainingMode enum, but got &#39;</span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span><span class="si">}</span><span class="s2">&#39;.&quot;</span>
        <span class="p">)</span>
    <span class="n">originally_training</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;training&quot;</span><span class="p">):</span>
        <span class="n">originally_training</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">training</span>

        <span class="c1"># ONNX opset 12 has better support for training amenable models, with updated</span>
        <span class="c1"># versions of the dropout and batch_norm operators</span>
        <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">TRAINING</span> <span class="ow">or</span> <span class="p">(</span>
            <span class="n">mode</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">PRESERVE</span> <span class="ow">and</span> <span class="n">originally_training</span>
        <span class="p">):</span>
            <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_training</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="k">if</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span> <span class="o">&lt;</span> <span class="mi">12</span><span class="p">:</span>
                <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
                    <span class="s2">&quot;You are exporting the model in training mode with onnx opset &quot;</span>
                    <span class="sa">f</span><span class="s2">&quot;version </span><span class="si">{</span><span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span><span class="si">}</span><span class="s2">. &quot;</span>
                    <span class="s2">&quot;Opset versions lower than opset 12 will not be able to export &quot;</span>
                    <span class="s2">&quot;nodes such as Dropout and BatchNorm correctly.&quot;</span>
                <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_training</span> <span class="o">=</span> <span class="kc">False</span>

        <span class="n">GLOBALS</span><span class="o">.</span><span class="n">training_mode</span> <span class="o">=</span> <span class="n">mode</span>
        <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">TRAINING</span><span class="p">:</span>
            <span class="n">model</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">mode</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span><span class="p">:</span>
            <span class="n">model</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
        <span class="c1"># else mode == _C_onnx.TrainingMode.PRESERVE, do nothing</span>

    <span class="k">try</span><span class="p">:</span>
        <span class="k">yield</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;training&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">mode</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">PRESERVE</span><span class="p">:</span>
            <span class="n">model</span><span class="o">.</span><span class="n">train</span><span class="p">(</span><span class="n">originally_training</span><span class="p">)</span></div>


<span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
<span class="k">def</span> <span class="nf">disable_apex_o2_state_dict_hook</span><span class="p">(</span><span class="n">model</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="o">|</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptFunction</span><span class="p">):</span>
    <span class="c1"># Apex O2 hook state_dict to return fp16 weights as fp32.</span>
    <span class="c1"># Exporter cannot identify them as same tensors.</span>
    <span class="c1"># Since this hook is only used by optimizer, it is safe to</span>
    <span class="c1"># remove this hook while exporting.</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptFunction</span><span class="p">):</span>
        <span class="n">model_hooks</span> <span class="o">=</span> <span class="p">{}</span>  <span class="c1"># type: ignore[var-annotated]</span>
        <span class="k">for</span> <span class="n">module</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">modules</span><span class="p">():</span>
            <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">hook</span> <span class="ow">in</span> <span class="n">module</span><span class="o">.</span><span class="n">_state_dict_hooks</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">hook</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span> <span class="o">==</span> <span class="s2">&quot;O2StateDictHook&quot;</span><span class="p">:</span>
                    <span class="k">if</span> <span class="n">module</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">model_hooks</span><span class="p">:</span>
                        <span class="n">model_hooks</span><span class="p">[</span><span class="n">module</span><span class="p">]</span> <span class="o">=</span> <span class="p">{}</span>
                    <span class="n">model_hooks</span><span class="p">[</span><span class="n">module</span><span class="p">][</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">hook</span>
            <span class="k">if</span> <span class="n">module</span> <span class="ow">in</span> <span class="n">model_hooks</span><span class="p">:</span>
                <span class="k">for</span> <span class="n">key</span> <span class="ow">in</span> <span class="n">model_hooks</span><span class="p">[</span><span class="n">module</span><span class="p">]:</span>
                    <span class="n">module</span><span class="o">.</span><span class="n">_state_dict_hooks</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="n">key</span><span class="p">)</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">yield</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="c1"># Add the hooks back</span>
            <span class="k">for</span> <span class="n">module</span><span class="p">,</span> <span class="n">m_map</span> <span class="ow">in</span> <span class="n">model_hooks</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">hook</span> <span class="ow">in</span> <span class="n">m_map</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                    <span class="n">module</span><span class="o">.</span><span class="n">_state_dict_hooks</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">hook</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="k">yield</span>
        <span class="k">finally</span><span class="p">:</span>
            <span class="k">pass</span>


<span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
<span class="k">def</span> <span class="nf">setup_onnx_logging</span><span class="p">(</span><span class="n">verbose</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
    <span class="n">is_originally_enabled</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_is_onnx_log_enabled</span>
    <span class="k">if</span> <span class="n">is_originally_enabled</span> <span class="ow">or</span> <span class="n">verbose</span><span class="p">:</span>  <span class="c1"># type: ignore[truthy-function]</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_set_onnx_log_enabled</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="k">yield</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">is_originally_enabled</span><span class="p">:</span>  <span class="c1"># type: ignore[truthy-function]</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_set_onnx_log_enabled</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>


<span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
<span class="k">def</span> <span class="nf">exporter_context</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="p">,</span> <span class="n">verbose</span><span class="p">:</span> <span class="nb">bool</span><span class="p">):</span>
    <span class="k">with</span> <span class="n">select_model_mode_for_export</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span> <span class="n">mode</span>
    <span class="p">)</span> <span class="k">as</span> <span class="n">mode_ctx</span><span class="p">,</span> <span class="n">disable_apex_o2_state_dict_hook</span><span class="p">(</span>
        <span class="n">model</span>
    <span class="p">)</span> <span class="k">as</span> <span class="n">apex_ctx</span><span class="p">,</span> <span class="n">setup_onnx_logging</span><span class="p">(</span>
        <span class="n">verbose</span>
    <span class="p">)</span> <span class="k">as</span> <span class="n">log_ctx</span><span class="p">,</span> <span class="n">diagnostics</span><span class="o">.</span><span class="n">create_export_diagnostic_context</span><span class="p">()</span> <span class="k">as</span> <span class="n">diagnostic_ctx</span><span class="p">:</span>
        <span class="k">yield</span> <span class="p">(</span><span class="n">mode_ctx</span><span class="p">,</span> <span class="n">apex_ctx</span><span class="p">,</span> <span class="n">log_ctx</span><span class="p">,</span> <span class="n">diagnostic_ctx</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_get_torch_export_args</span><span class="p">(</span>
    <span class="n">args</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span>
    <span class="n">kwargs</span><span class="p">:</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span><span class="nb">tuple</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span><span class="p">]:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Obtain the arguments for torch.onnx.export from the model and the input arguments.&quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">kwargs</span> <span class="ow">and</span> <span class="n">args</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">kwargs</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">args</span> <span class="o">=</span> <span class="n">args</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span>


<span class="k">def</span> <span class="nf">export</span><span class="p">(</span>
    <span class="n">model</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="o">|</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span> <span class="o">|</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptFunction</span><span class="p">,</span>
    <span class="n">args</span><span class="p">:</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span> <span class="o">|</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span>
    <span class="n">f</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
    <span class="o">*</span><span class="p">,</span>
    <span class="n">kwargs</span><span class="p">:</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">export_params</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
    <span class="n">verbose</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="n">training</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span> <span class="o">=</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span><span class="p">,</span>
    <span class="n">input_names</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">output_names</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">operator_export_type</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span> <span class="o">=</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span><span class="p">,</span>
    <span class="n">opset_version</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">do_constant_folding</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
    <span class="n">dynamic_axes</span><span class="p">:</span> <span class="n">Mapping</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Mapping</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">str</span><span class="p">]]</span>
    <span class="o">|</span> <span class="n">Mapping</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span>
    <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">keep_initializers_as_inputs</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">custom_opsets</span><span class="p">:</span> <span class="n">Mapping</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="nb">int</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">export_modules_as_functions</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">|</span> <span class="n">Collection</span><span class="p">[</span><span class="nb">type</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="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="n">autograd_inlining</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Exports a model into ONNX format.</span>

<span class="sd">    If ``model`` is not a :class:`torch.jit.ScriptModule` nor a</span>
<span class="sd">    :class:`torch.jit.ScriptFunction`, this runs</span>
<span class="sd">    ``model`` once in order to convert it to a TorchScript graph to be exported</span>
<span class="sd">    (the equivalent of :func:`torch.jit.trace`). Thus this has the same limited support</span>
<span class="sd">    for dynamic control flow as :func:`torch.jit.trace`.</span>

<span class="sd">    Args:</span>
<span class="sd">        model: The model to be exported.</span>
<span class="sd">        args:</span>

<span class="sd">            args can be structured either as:</span>

<span class="sd">            1. ONLY A TUPLE OF ARGUMENTS::</span>

<span class="sd">                args = (x, y, z)</span>

<span class="sd">            The tuple should contain model inputs such that ``model(*args)`` is a valid</span>
<span class="sd">            invocation of the model. Any non-Tensor arguments will be hard-coded into the</span>
<span class="sd">            exported model; any Tensor arguments will become inputs of the exported model,</span>
<span class="sd">            in the order they occur in the tuple.</span>

<span class="sd">            2. A TENSOR::</span>

<span class="sd">                args = torch.Tensor([1])</span>

<span class="sd">            This is equivalent to a 1-ary tuple of that Tensor.</span>

<span class="sd">            3. A TUPLE OF ARGUMENTS ENDING WITH A DICTIONARY OF NAMED ARGUMENTS::</span>

<span class="sd">                args = (x, {&quot;y&quot;: input_y, &quot;z&quot;: input_z})</span>

<span class="sd">            All but the last element of the tuple will be passed as non-keyword arguments,</span>
<span class="sd">            and named arguments will be set from the last element. If a named argument is</span>
<span class="sd">            not present in the dictionary, it is assigned the default value, or None if a</span>
<span class="sd">            default value is not provided.</span>

<span class="sd">            .. warning::</span>
<span class="sd">                This behavior will be deprecated in a future release. Please use the</span>
<span class="sd">                kwargs argument instead.</span>

<span class="sd">            .. note::</span>
<span class="sd">                If a dictionary is the last element of the args tuple, it will be</span>
<span class="sd">                interpreted as containing named arguments. In order to pass a dict as the</span>
<span class="sd">                last non-keyword arg, provide an empty dict as the last element of the args</span>
<span class="sd">                tuple. For example, instead of::</span>

<span class="sd">                    torch.onnx.export(</span>
<span class="sd">                        model,</span>
<span class="sd">                        (</span>
<span class="sd">                            x,</span>
<span class="sd">                            # WRONG: will be interpreted as named arguments</span>
<span class="sd">                            {y: z},</span>
<span class="sd">                        ),</span>
<span class="sd">                        &quot;test.onnx.pb&quot;,</span>
<span class="sd">                    )</span>

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

<span class="sd">                    torch.onnx.export(model, (x, {y: z}, {}), &quot;test.onnx.pb&quot;)</span>

<span class="sd">        f: Path to the output ONNX model file. E.g. &quot;model.onnx&quot;.</span>
<span class="sd">        kwargs: Named arguments to the model.</span>
<span class="sd">        export_params: If True, all parameters will</span>
<span class="sd">            be exported. Set this to False if you want to export an untrained model.</span>
<span class="sd">            In this case, the exported model will first take all of its parameters</span>
<span class="sd">            as arguments, with the ordering as specified by ``model.state_dict().values()``</span>
<span class="sd">        verbose: if True, prints a description of the</span>
<span class="sd">            model being exported to stdout. In addition, the final ONNX graph will include the</span>
<span class="sd">            field ``doc_string``` from the exported model which mentions the source code locations</span>
<span class="sd">            for ``model``. If True, ONNX exporter logging will be turned on.</span>
<span class="sd">        training:</span>
<span class="sd">            * ``TrainingMode.EVAL``: export the model in inference mode.</span>
<span class="sd">            * ``TrainingMode.PRESERVE``: export the model in inference mode if model.training is</span>
<span class="sd">                False and in training mode if model.training is True.</span>
<span class="sd">            * ``TrainingMode.TRAINING``: export the model in training mode. Disables optimizations</span>
<span class="sd">                which might interfere with training.</span>
<span class="sd">        input_names (list of str, default empty list): names to assign to the</span>
<span class="sd">            input nodes of the graph, in order.</span>
<span class="sd">        output_names (list of str, default empty list): names to assign to the</span>
<span class="sd">            output nodes of the graph, in order.</span>
<span class="sd">        operator_export_type (enum, default OperatorExportTypes.ONNX):</span>

<span class="sd">            .. warning::</span>
<span class="sd">                This option will be deprecated in a future release. Future exported</span>
<span class="sd">                graphs will always use the default opset domain.</span>

<span class="sd">            * ``OperatorExportTypes.ONNX``: Export all ops as regular ONNX ops</span>
<span class="sd">                (in the default opset domain).</span>
<span class="sd">            * ``OperatorExportTypes.ONNX_FALLTHROUGH``: Try to convert all ops</span>
<span class="sd">                to standard ONNX ops in the default opset domain. If unable to do so</span>
<span class="sd">                (e.g. because support has not been added to convert a particular torch op to ONNX),</span>
<span class="sd">                fall back to exporting the op into a custom opset domain without conversion. Applies</span>
<span class="sd">                to `custom ops &lt;https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html&gt;`_</span>
<span class="sd">                as well as ATen ops. For the exported model to be usable, the runtime must support</span>
<span class="sd">                these non-standard ops.</span>
<span class="sd">            * ``OperatorExportTypes.ONNX_ATEN``: All ATen ops (in the TorchScript namespace &quot;aten&quot;)</span>
<span class="sd">                are exported as ATen ops (in opset domain &quot;org.pytorch.aten&quot;).</span>
<span class="sd">                `ATen &lt;https://pytorch.org/cppdocs/#aten&gt;`_ is PyTorch&#39;s built-in tensor library, so</span>
<span class="sd">                this instructs the runtime to use PyTorch&#39;s implementation of these ops.</span>

<span class="sd">                .. warning::</span>

<span class="sd">                    Models exported this way are probably runnable only by Caffe2.</span>

<span class="sd">                    This may be useful if the numeric differences in implementations of operators are</span>
<span class="sd">                    causing large differences in behavior between PyTorch and Caffe2 (which is more</span>
<span class="sd">                    common on untrained models).</span>

<span class="sd">            * ``OperatorExportTypes.ONNX_ATEN_FALLBACK``: Try to export each ATen op</span>
<span class="sd">                (in the TorchScript namespace &quot;aten&quot;) as a regular ONNX op. If we are unable to do so</span>
<span class="sd">                (e.g. because support has not been added to convert a particular torch op to ONNX),</span>
<span class="sd">                fall back to exporting an ATen op. See documentation on OperatorExportTypes.ONNX_ATEN for</span>
<span class="sd">                context.</span>
<span class="sd">                For example::</span>

<span class="sd">                    graph(%0 : Float):</span>
<span class="sd">                    %3 : int = prim::Constant[value=0]()</span>
<span class="sd">                    # conversion unsupported</span>
<span class="sd">                    %4 : Float = aten::triu(%0, %3)</span>
<span class="sd">                    # conversion supported</span>
<span class="sd">                    %5 : Float = aten::mul(%4, %0)</span>
<span class="sd">                    return (%5)</span>

<span class="sd">                Assuming ``aten::triu`` is not supported in ONNX, this will be exported as::</span>

<span class="sd">                    graph(%0 : Float):</span>
<span class="sd">                    %1 : Long() = onnx::Constant[value={0}]()</span>
<span class="sd">                    # not converted</span>
<span class="sd">                    %2 : Float = aten::ATen[operator=&quot;triu&quot;](%0, %1)</span>
<span class="sd">                    # converted</span>
<span class="sd">                    %3 : Float = onnx::Mul(%2, %0)</span>
<span class="sd">                    return (%3)</span>

<span class="sd">                .. warning::</span>

<span class="sd">                    Models exported this way are probably runnable only by Caffe2.</span>

<span class="sd">        opset_version (int, default 17): The version of the</span>
<span class="sd">            `default (ai.onnx) opset &lt;https://github.com/onnx/onnx/blob/master/docs/Operators.md&gt;`_</span>
<span class="sd">            to target. Must be &gt;= 7 and &lt;= 17.</span>
<span class="sd">        do_constant_folding: Apply the constant-folding optimization.</span>
<span class="sd">            Constant-folding will replace some of the ops that have all constant inputs</span>
<span class="sd">            with pre-computed constant nodes.</span>
<span class="sd">        dynamic_axes:</span>

<span class="sd">            By default the exported model will have the shapes of all input and output tensors</span>
<span class="sd">            set to exactly match those given in ``args``. To specify axes of tensors as</span>
<span class="sd">            dynamic (i.e. known only at run-time), set ``dynamic_axes`` to a dict with schema:</span>

<span class="sd">            * KEY (str): an input or output name. Each name must also be provided in ``input_names`` or</span>
<span class="sd">                ``output_names``.</span>
<span class="sd">            * VALUE (dict or list): If a dict, keys are axis indices and values are axis names. If a</span>
<span class="sd">                list, each element is an axis index.</span>

<span class="sd">            For example::</span>

<span class="sd">                class SumModule(torch.nn.Module):</span>
<span class="sd">                    def forward(self, x):</span>
<span class="sd">                        return torch.sum(x, dim=1)</span>


<span class="sd">                torch.onnx.export(</span>
<span class="sd">                    SumModule(),</span>
<span class="sd">                    (torch.ones(2, 2),),</span>
<span class="sd">                    &quot;onnx.pb&quot;,</span>
<span class="sd">                    input_names=[&quot;x&quot;],</span>
<span class="sd">                    output_names=[&quot;sum&quot;],</span>
<span class="sd">                )</span>

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

<span class="sd">                input {</span>
<span class="sd">                  name: &quot;x&quot;</span>
<span class="sd">                  ...</span>
<span class="sd">                      shape {</span>
<span class="sd">                        dim {</span>
<span class="sd">                          dim_value: 2  # axis 0</span>
<span class="sd">                        }</span>
<span class="sd">                        dim {</span>
<span class="sd">                          dim_value: 2  # axis 1</span>
<span class="sd">                ...</span>
<span class="sd">                output {</span>
<span class="sd">                  name: &quot;sum&quot;</span>
<span class="sd">                  ...</span>
<span class="sd">                      shape {</span>
<span class="sd">                        dim {</span>
<span class="sd">                          dim_value: 2  # axis 0</span>
<span class="sd">                ...</span>

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

<span class="sd">                torch.onnx.export(</span>
<span class="sd">                    SumModule(),</span>
<span class="sd">                    (torch.ones(2, 2),),</span>
<span class="sd">                    &quot;onnx.pb&quot;,</span>
<span class="sd">                    input_names=[&quot;x&quot;],</span>
<span class="sd">                    output_names=[&quot;sum&quot;],</span>
<span class="sd">                    dynamic_axes={</span>
<span class="sd">                        # dict value: manually named axes</span>
<span class="sd">                        &quot;x&quot;: {0: &quot;my_custom_axis_name&quot;},</span>
<span class="sd">                        # list value: automatic names</span>
<span class="sd">                        &quot;sum&quot;: [0],</span>
<span class="sd">                    },</span>
<span class="sd">                )</span>

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

<span class="sd">                input {</span>
<span class="sd">                  name: &quot;x&quot;</span>
<span class="sd">                  ...</span>
<span class="sd">                      shape {</span>
<span class="sd">                        dim {</span>
<span class="sd">                          dim_param: &quot;my_custom_axis_name&quot;  # axis 0</span>
<span class="sd">                        }</span>
<span class="sd">                        dim {</span>
<span class="sd">                          dim_value: 2  # axis 1</span>
<span class="sd">                ...</span>
<span class="sd">                output {</span>
<span class="sd">                  name: &quot;sum&quot;</span>
<span class="sd">                  ...</span>
<span class="sd">                      shape {</span>
<span class="sd">                        dim {</span>
<span class="sd">                          dim_param: &quot;sum_dynamic_axes_1&quot;  # axis 0</span>
<span class="sd">                ...</span>

<span class="sd">        keep_initializers_as_inputs: If True, all the</span>
<span class="sd">            initializers (typically corresponding to parameters) in the</span>
<span class="sd">            exported graph will also be added as inputs to the graph. If False,</span>
<span class="sd">            then initializers are not added as inputs to the graph, and only</span>
<span class="sd">            the non-parameter inputs are added as inputs.</span>
<span class="sd">            This may allow for better optimizations (e.g. constant folding) by</span>
<span class="sd">            backends/runtimes.</span>

<span class="sd">            If True, `deduplicate_initializers` pass will not be executed. This means</span>
<span class="sd">            initializers with duplicated values will not be deduplicated and</span>
<span class="sd">            will be treated as distinct inputs to the graph. This allows different</span>
<span class="sd">            input initializers to be supplied at the runtime following export.</span>

<span class="sd">            If ``opset_version &lt; 9``, initializers MUST be part of graph</span>
<span class="sd">            inputs and this argument will be ignored and the behavior will be</span>
<span class="sd">            equivalent to setting this argument to True.</span>

<span class="sd">        custom_opsets (dict[str, int], default empty dict): A dict with schema:</span>

<span class="sd">            * KEY (str): opset domain name</span>
<span class="sd">            * VALUE (int): opset version</span>

<span class="sd">            If a custom opset is referenced by ``model`` but not mentioned in this dictionary,</span>
<span class="sd">            the opset version is set to 1. Only custom opset domain name and version should be</span>
<span class="sd">            indicated through this argument.</span>

<span class="sd">        export_modules_as_functions: Flag to enable</span>
<span class="sd">            exporting all ``nn.Module`` forward calls as local functions in ONNX. Or a set to indicate the</span>
<span class="sd">            particular types of modules to export as local functions in ONNX.</span>
<span class="sd">            This feature requires ``opset_version`` &gt;= 15, otherwise the export will fail. This is because</span>
<span class="sd">            ``opset_version`` &lt; 15 implies IR version &lt; 8, which means no local function support.</span>
<span class="sd">            Module variables will be exported as function attributes. There are two categories of function</span>
<span class="sd">            attributes.</span>

<span class="sd">            1. Annotated attributes: class variables that have type annotations via</span>
<span class="sd">            `PEP 526-style &lt;https://www.python.org/dev/peps/pep-0526/#class-and-instance-variable-annotations&gt;`_</span>
<span class="sd">            will be exported as attributes.</span>
<span class="sd">            Annotated attributes are not used inside the subgraph of ONNX local function because</span>
<span class="sd">            they are not created by PyTorch JIT tracing, but they may be used by consumers</span>
<span class="sd">            to determine whether or not to replace the function with a particular fused kernel.</span>

<span class="sd">            2. Inferred attributes: variables that are used by operators inside the module. Attribute names</span>
<span class="sd">            will have prefix &quot;inferred::&quot;. This is to differentiate from predefined attributes retrieved from</span>
<span class="sd">            python module annotations. Inferred attributes are used inside the subgraph of ONNX local function.</span>

<span class="sd">            * ``False`` (default): export ``nn.Module`` forward calls as fine grained nodes.</span>
<span class="sd">            * ``True``: export all ``nn.Module`` forward calls as local function nodes.</span>
<span class="sd">            * Set of type of nn.Module: export ``nn.Module`` forward calls as local function nodes,</span>
<span class="sd">                only if the type of the ``nn.Module`` is found in the set.</span>

<span class="sd">        autograd_inlining: Flag used to control whether to inline autograd functions.</span>
<span class="sd">            Refer to https://github.com/pytorch/pytorch/pull/74765 for more details.</span>

<span class="sd">    Raises:</span>
<span class="sd">        :class:`torch.onnx.errors.CheckerError`: If the ONNX checker detects an invalid ONNX graph.</span>
<span class="sd">        :class:`torch.onnx.errors.UnsupportedOperatorError`: If the ONNX graph cannot be exported because it</span>
<span class="sd">            uses an operator that is not supported by the exporter.</span>
<span class="sd">        :class:`torch.onnx.errors.OnnxExporterError`: Other errors that can occur during export.</span>
<span class="sd">            All errors are subclasses of :class:`errors.OnnxExporterError`.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">operator_export_type</span> <span class="o">!=</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span><span class="p">:</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
            <span class="s2">&quot;Setting `operator_export_type` to something other than default is deprecated. &quot;</span>
            <span class="s2">&quot;The option will be removed in a future release.&quot;</span><span class="p">,</span>
            <span class="n">category</span><span class="o">=</span><span class="ne">FutureWarning</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">training</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">TRAINING</span><span class="p">:</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
            <span class="s2">&quot;Setting `training` to something other than default is deprecated. &quot;</span>
            <span class="s2">&quot;The option will be removed in a future release. Please set the training mode &quot;</span>
            <span class="s2">&quot;before exporting the model.&quot;</span><span class="p">,</span>
            <span class="n">category</span><span class="o">=</span><span class="ne">FutureWarning</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="n">args</span><span class="p">,)</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="k">else</span> <span class="n">args</span>
    <span class="k">if</span> <span class="n">kwargs</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">args</span> <span class="o">=</span> <span class="n">args</span> <span class="o">+</span> <span class="p">(</span><span class="n">kwargs</span><span class="p">,)</span>

    <span class="n">_export</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">args</span><span class="p">,</span>
        <span class="n">f</span><span class="p">,</span>
        <span class="n">export_params</span><span class="p">,</span>
        <span class="n">verbose</span><span class="p">,</span>
        <span class="n">training</span><span class="p">,</span>
        <span class="n">input_names</span><span class="p">,</span>
        <span class="n">output_names</span><span class="p">,</span>
        <span class="n">operator_export_type</span><span class="o">=</span><span class="n">operator_export_type</span><span class="p">,</span>
        <span class="n">opset_version</span><span class="o">=</span><span class="n">opset_version</span><span class="p">,</span>
        <span class="n">do_constant_folding</span><span class="o">=</span><span class="n">do_constant_folding</span><span class="p">,</span>
        <span class="n">dynamic_axes</span><span class="o">=</span><span class="n">dynamic_axes</span><span class="p">,</span>
        <span class="n">keep_initializers_as_inputs</span><span class="o">=</span><span class="n">keep_initializers_as_inputs</span><span class="p">,</span>
        <span class="n">custom_opsets</span><span class="o">=</span><span class="n">custom_opsets</span><span class="p">,</span>
        <span class="n">export_modules_as_functions</span><span class="o">=</span><span class="n">export_modules_as_functions</span><span class="p">,</span>
        <span class="n">autograd_inlining</span><span class="o">=</span><span class="n">autograd_inlining</span><span class="p">,</span>
    <span class="p">)</span>

    <span class="k">return</span> <span class="kc">None</span>


<span class="k">def</span> <span class="nf">_is_constant_tensor_list</span><span class="p">(</span><span class="n">node</span><span class="p">):</span>
    <span class="k">if</span> <span class="n">node</span><span class="o">.</span><span class="n">kind</span><span class="p">()</span> <span class="o">!=</span> <span class="s2">&quot;prim::Constant&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="kc">False</span>
    <span class="n">output_type</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">output</span><span class="p">()</span><span class="o">.</span><span class="n">type</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">output_type</span><span class="o">.</span><span class="n">isSubtypeOf</span><span class="p">(</span><span class="n">_C</span><span class="o">.</span><span class="n">ListType</span><span class="o">.</span><span class="n">ofTensors</span><span class="p">()):</span>
        <span class="k">return</span> <span class="kc">True</span>
    <span class="k">if</span> <span class="n">output_type</span><span class="o">.</span><span class="n">isSubtypeOf</span><span class="p">(</span><span class="n">_C</span><span class="o">.</span><span class="n">ListType</span><span class="p">(</span><span class="n">_C</span><span class="o">.</span><span class="n">OptionalType</span><span class="o">.</span><span class="n">ofTensor</span><span class="p">())):</span>
        <span class="k">return</span> <span class="kc">True</span>


<span class="c1"># ONNX can&#39;t handle constants that are lists of tensors, which can</span>
<span class="c1"># get generated in constant prop. So we split them back into prim::ListConstructs</span>


<span class="k">def</span> <span class="nf">_split_tensor_list_constants</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">block</span><span class="p">):</span>
    <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">block</span><span class="o">.</span><span class="n">nodes</span><span class="p">():</span>
        <span class="k">for</span> <span class="n">subblock</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">blocks</span><span class="p">():</span>
            <span class="n">_split_tensor_list_constants</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">subblock</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">_is_constant_tensor_list</span><span class="p">(</span><span class="n">node</span><span class="p">):</span>
            <span class="n">inputs</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="k">for</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">output</span><span class="p">()</span><span class="o">.</span><span class="n">toIValue</span><span class="p">():</span>
                <span class="nb">input</span> <span class="o">=</span> <span class="n">g</span><span class="o">.</span><span class="n">insertConstant</span><span class="p">(</span><span class="n">val</span><span class="p">)</span>
                <span class="nb">input</span><span class="o">.</span><span class="n">node</span><span class="p">()</span><span class="o">.</span><span class="n">moveBefore</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>
                <span class="nb">input</span><span class="o">.</span><span class="n">node</span><span class="p">()</span><span class="o">.</span><span class="n">copyMetadata</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>
                <span class="n">inputs</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>

            <span class="n">lc</span> <span class="o">=</span> <span class="p">(</span>
                <span class="n">g</span><span class="o">.</span><span class="n">create</span><span class="p">(</span><span class="s2">&quot;prim::ListConstruct&quot;</span><span class="p">,</span> <span class="n">inputs</span><span class="p">)</span>
                <span class="o">.</span><span class="n">insertBefore</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>
                <span class="o">.</span><span class="n">output</span><span class="p">()</span>
                <span class="o">.</span><span class="n">setType</span><span class="p">(</span><span class="n">_C</span><span class="o">.</span><span class="n">ListType</span><span class="o">.</span><span class="n">ofTensors</span><span class="p">())</span>
            <span class="p">)</span>
            <span class="n">lc</span><span class="o">.</span><span class="n">node</span><span class="p">()</span><span class="o">.</span><span class="n">copyMetadata</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>
            <span class="n">node</span><span class="o">.</span><span class="n">output</span><span class="p">()</span><span class="o">.</span><span class="n">replaceAllUsesWith</span><span class="p">(</span><span class="n">lc</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_optimize_graph</span><span class="p">(</span>
    <span class="n">graph</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Graph</span><span class="p">,</span>
    <span class="n">operator_export_type</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="p">,</span>
    <span class="n">_disable_torch_constant_prop</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="n">fixed_batch_size</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="n">params_dict</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">dynamic_axes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">input_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">module</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
<span class="p">):</span>
    <span class="k">if</span> <span class="n">params_dict</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">params_dict</span> <span class="o">=</span> <span class="p">{}</span>

    <span class="c1"># Inline everything</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_inline</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># Remove fork/wait nodes</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_inline_fork_wait</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">autograd_inlining</span><span class="p">:</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_autograd_function_process</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lower_all_tuples</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># we now record some ops like ones/zeros</span>
    <span class="c1"># into a trace where we previously recorded constants.</span>
    <span class="c1"># use constant prop to maintain our current level of onnx support</span>
    <span class="c1"># without implementing symbolics for all of them</span>
    <span class="k">if</span> <span class="n">_disable_torch_constant_prop</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_constant_propagation</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">_split_tensor_list_constants</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">graph</span><span class="p">)</span>
    <span class="c1"># run dce to eliminate dead parts of the graph that might have been</span>
    <span class="c1"># left behind by things like symbolic_override</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_dce</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># CSE should improve perf when Autocast is used with disabled cache</span>
    <span class="c1"># Autocast is disabled due to a limitation on tracer as described at https://github.com/pytorch/pytorch/issues/84092</span>
    <span class="c1"># Must run before _C._jit_pass_erase_number_types to prevent type substitution</span>
    <span class="k">if</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_cse</span><span class="p">(</span><span class="n">graph</span><span class="p">):</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_canonicalize_graph_fuser_ops</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_peephole</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_fuse_addmm</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_peephole</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lower_all_tuples</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="c1"># in _jit_pass_onnx, symbolic functions are called for each node for conversion.</span>
    <span class="c1"># However, there are nodes that cannot be converted without additional context.</span>
    <span class="c1"># For example, the number of outputs from split (and whether it is static or dynamic) is unknown</span>
    <span class="c1"># until the point where it is unpacked by listUnpack node.</span>
    <span class="c1"># This pass does a preprocess, and prepares the nodes such that enough context can be received</span>
    <span class="c1"># by the symbolic function.</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_remove_inplace_ops_for_onnx</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">module</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_preprocess</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># onnx does not support tuples, so try to remove them</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># onnx only supports tensors, but 1 / 2 = 0.5 and tensor(1) / tensor(2) = 0</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_prepare_division_for_onnx</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_remove_print</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_preprocess_caffe2</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">symbolic_helper</span><span class="o">.</span><span class="n">_quantized_ops</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
    <span class="c1"># Unpack quantized weights for conv and linear ops and insert into graph.</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_unpack_quantized_weights</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">)</span>
    <span class="c1"># onnx only supports tensors, so we turn all out number types into tensors</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_erase_number_types</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">onnx_shape_inference</span><span class="p">:</span>
        <span class="n">input_names</span> <span class="o">=</span> <span class="p">[]</span> <span class="k">if</span> <span class="n">input_names</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">input_names</span>
        <span class="n">dynamic_axes</span> <span class="o">=</span> <span class="p">{}</span> <span class="k">if</span> <span class="n">dynamic_axes</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">dynamic_axes</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_set_dynamic_input_shape</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">dynamic_axes</span><span class="p">,</span> <span class="n">input_names</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">graph</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_scalar_type_analysis</span><span class="p">(</span>
        <span class="n">graph</span><span class="p">,</span> <span class="kc">True</span><span class="p">,</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span>
    <span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_peephole</span><span class="p">(</span>
        <span class="n">graph</span><span class="p">,</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span><span class="p">,</span> <span class="n">fixed_batch_size</span>
    <span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># graph is not a valid jit graph anymore because types have been replaced</span>
    <span class="c1"># (e.g. int with Tensor), so it now contains operators that don&#39;t actually</span>
    <span class="c1"># exist. We can&#39;t run normal dead code elimination because it&#39;d fail trying</span>
    <span class="c1"># to look up if an operator has side effects, but we can run a dead code</span>
    <span class="c1"># elimination variant that doesn&#39;t need to look up if an op has side effects.</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_dce_allow_deleting_nodes_with_side_effects</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">graph</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_canonicalize</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">onnx_shape_inference</span><span class="p">:</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_graph_shape_type_inference</span><span class="p">(</span>
                <span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">,</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span>
            <span class="p">)</span>
        <span class="k">except</span> <span class="ne">RuntimeError</span><span class="p">:</span>
            <span class="c1"># NOTE: shape type inference error should not stop the export process</span>
            <span class="c1"># https://github.com/pytorch/pytorch/issues/132205</span>
            <span class="k">pass</span>

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


<span class="k">def</span> <span class="nf">warn_on_static_input_change</span><span class="p">(</span><span class="n">input_states</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Warns that changes to input dictionaries and strings won&#39;t take effect in the traced ONNX graph.</span>

<span class="sd">    We accept dictionaries and strings as ONNX inputs, but they should be only for</span>
<span class="sd">    configuration use. we detect here if these inputs are modified, and if so we warn</span>
<span class="sd">    the user that the changes won&#39;t take effect in the traced ONNX graph.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">traced_input</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">input_states</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">input_states</span><span class="p">[</span><span class="mi">1</span><span class="p">]):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">list</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> <span class="o">!=</span> <span class="nb">list</span><span class="p">(</span><span class="n">traced_input</span><span class="o">.</span><span class="n">keys</span><span class="p">()):</span>
                <span class="n">warning</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="s2">&quot;We detected that you are modifying a dictionary that is an input to your &quot;</span>
                    <span class="s2">&quot;model. &quot;</span>
                    <span class="s2">&quot;Note that dictionaries are allowed as inputs in ONNX but they should be &quot;</span>
                    <span class="s2">&quot;handled with care. &quot;</span>
                    <span class="s2">&quot;Usages of dictionaries is not recommended, and should not be used except &quot;</span>
                    <span class="s2">&quot;for configuration use. &quot;</span>
                    <span class="s2">&quot;Also note that the order and values of the keys must remain the same. &quot;</span>
                <span class="p">)</span>
                <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">warning</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">input</span> <span class="o">!=</span> <span class="n">traced_input</span><span class="p">:</span>
                <span class="n">warning</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="s2">&quot;The model seems to have string inputs/outputs. &quot;</span>
                    <span class="s2">&quot;Note that strings will not appear as inputs/outputs of the ONNX graph. &quot;</span>
                <span class="p">)</span>
                <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">warning</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_resolve_args_by_export_type</span><span class="p">(</span><span class="n">arg_name</span><span class="p">,</span> <span class="n">arg_value</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Resolves the arguments that are ignored when export_type != operator_export_type.ONNX.&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">arg_value</span>


<span class="k">def</span> <span class="nf">_decide_keep_init_as_input</span><span class="p">(</span>
    <span class="n">keep_initializers_as_inputs</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">|</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">operator_export_type</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="p">,</span>
    <span class="n">opset_version</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Decides whether the initializers in the graph should be listed as ONNX graph inputs.</span>

<span class="sd">    This method encapsulates the logic to decide whether the initializers in the graph</span>
<span class="sd">    should be listed as ONNX graph inputs (i.e., whether to choose ONNX IR v3 or v4).</span>
<span class="sd">    If keep_initializers_as_inputs is not specified (None), then we decide whether to keep</span>
<span class="sd">    initializers as graph inputs (val_keep_init_as_ip) based on export type. If export type</span>
<span class="sd">    is ONNX, then do not keep initializers as input (val_keep_init_as_ip=False). For all other</span>
<span class="sd">    export types keep initializers as input (val_keep_init_as_ip=True).</span>
<span class="sd">    If keep_initializers_as_inputs is specified, then respect it. Unless opset version &lt;= 8,</span>
<span class="sd">    in which case it must be ignored because for opset version &lt;= 8, all initializers MUST be</span>
<span class="sd">    part of graph input (only ONNX IR v3 is allowed), i.e. val_keep_init_as_ip=True.</span>

<span class="sd">    Special handling is needed for opset version 8 or lower, because irrespective</span>
<span class="sd">    of user input for keep_initializers_as_inputs, the graph must follow ONNX IR v3</span>
<span class="sd">    semantics, i.e. all initializers must be listed as ONNX graph input.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">opset_version</span> <span class="o">&lt;</span> <span class="mi">9</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">keep_initializers_as_inputs</span> <span class="ow">is</span> <span class="kc">False</span><span class="p">:</span>
            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
                <span class="s2">&quot;Setting &#39;keep_initializers_as_inputs=False&#39; for opset version&quot;</span>
                <span class="s2">&quot;8 or lower would lead to an invalid ONNX graph. Therefore, &quot;</span>
                <span class="s2">&quot;&#39;keep_initializers_as_inputs=False&#39; is ignored during export.&quot;</span>
                <span class="s2">&quot;Exported model will have initializers as graph inputs (compliant &quot;</span>
                <span class="s2">&quot; to ONNX IR v3).&quot;</span>
            <span class="p">)</span>
        <span class="k">return</span> <span class="kc">True</span>  <span class="c1"># i.e. True == initializers are part of graph input (ONNX IR v3)</span>
    <span class="n">val_keep_init_as_ip</span> <span class="o">=</span> <span class="p">(</span>
        <span class="kc">True</span> <span class="k">if</span> <span class="n">keep_initializers_as_inputs</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">keep_initializers_as_inputs</span>
    <span class="p">)</span>
    <span class="k">if</span> <span class="p">(</span>
        <span class="n">keep_initializers_as_inputs</span> <span class="ow">is</span> <span class="kc">None</span>
        <span class="ow">and</span> <span class="n">operator_export_type</span> <span class="ow">is</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span>
    <span class="p">):</span>
        <span class="n">val_keep_init_as_ip</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="k">return</span> <span class="n">val_keep_init_as_ip</span>


<span class="k">def</span> <span class="nf">_decide_add_node_names</span><span class="p">(</span><span class="n">add_node_names</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">_resolve_args_by_export_type</span><span class="p">(</span>
        <span class="s2">&quot;add_node_names&quot;</span><span class="p">,</span> <span class="n">add_node_names</span><span class="p">,</span> <span class="n">operator_export_type</span>
    <span class="p">)</span>


<span class="k">def</span> <span class="nf">_decide_constant_folding</span><span class="p">(</span><span class="n">do_constant_folding</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">,</span> <span class="n">training</span><span class="p">):</span>
    <span class="n">do_constant_folding</span> <span class="o">=</span> <span class="n">_resolve_args_by_export_type</span><span class="p">(</span>
        <span class="s2">&quot;do_constant_folding&quot;</span><span class="p">,</span> <span class="n">do_constant_folding</span><span class="p">,</span> <span class="n">operator_export_type</span>
    <span class="p">)</span>
    <span class="k">if</span> <span class="n">do_constant_folding</span> <span class="ow">and</span> <span class="p">(</span>
        <span class="n">training</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">training</span> <span class="ow">is</span> <span class="ow">not</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span>
    <span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
            <span class="s2">&quot;It is recommended that constant folding be turned off (&#39;do_constant_folding=False&#39;) &quot;</span>
            <span class="s2">&quot;when exporting the model in training-amenable mode, i.e. with &#39;training=TrainingMode.TRAIN&#39; &quot;</span>
            <span class="s2">&quot;or &#39;training=TrainingMode.PRESERVE&#39; (when model is in training mode). Otherwise, some &quot;</span>
            <span class="s2">&quot;learnable model parameters may not translate correctly in the exported ONNX model &quot;</span>
            <span class="s2">&quot;because constant folding mutates model parameters. Please consider &quot;</span>
            <span class="s2">&quot;turning off constant folding or setting the training=TrainingMode.EVAL.&quot;</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">do_constant_folding</span>


<span class="k">def</span> <span class="nf">_signature</span><span class="p">(</span><span class="n">model</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">inspect</span><span class="o">.</span><span class="n">Signature</span><span class="p">:</span>
    <span class="n">should_be_callable</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;forward&quot;</span><span class="p">,</span> <span class="n">model</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">callable</span><span class="p">(</span><span class="n">should_be_callable</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">inspect</span><span class="o">.</span><span class="n">signature</span><span class="p">(</span><span class="n">should_be_callable</span><span class="p">)</span>
    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;model has no forward method and is not callable&quot;</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_decide_input_format</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">):</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">sig</span> <span class="o">=</span> <span class="n">_signature</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
    <span class="k">except</span> <span class="ne">ValueError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s2">, skipping _decide_input_format&quot;</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">args</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">ordered_list_keys</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">sig</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
        <span class="k">if</span> <span class="n">ordered_list_keys</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="s2">&quot;self&quot;</span><span class="p">:</span>
            <span class="n">ordered_list_keys</span> <span class="o">=</span> <span class="n">ordered_list_keys</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>
        <span class="n">args_dict</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">args_list</span> <span class="o">=</span> <span class="n">args</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="n">args_list</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">args</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">args_list</span> <span class="o">=</span> <span class="p">[</span><span class="n">args</span><span class="p">]</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args_list</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="n">args_dict</span> <span class="o">=</span> <span class="n">args_list</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
            <span class="n">args_list</span> <span class="o">=</span> <span class="n">args_list</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">n_nonkeyword</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">args_list</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">optional_arg</span> <span class="ow">in</span> <span class="n">ordered_list_keys</span><span class="p">[</span><span class="n">n_nonkeyword</span><span class="p">:]:</span>
            <span class="k">if</span> <span class="n">optional_arg</span> <span class="ow">in</span> <span class="n">args_dict</span><span class="p">:</span>
                <span class="n">args_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">args_dict</span><span class="p">[</span><span class="n">optional_arg</span><span class="p">])</span>
            <span class="c1"># Check if this arg has a default value</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">param</span> <span class="o">=</span> <span class="n">sig</span><span class="o">.</span><span class="n">parameters</span><span class="p">[</span><span class="n">optional_arg</span><span class="p">]</span>
                <span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">default</span> <span class="o">!=</span> <span class="n">param</span><span class="o">.</span><span class="n">empty</span><span class="p">:</span>
                    <span class="n">args_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">param</span><span class="o">.</span><span class="n">default</span><span class="p">)</span>
        <span class="n">args</span> <span class="o">=</span> <span class="n">args_list</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="nb">list</span><span class="p">)</span> <span class="k">else</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">args_list</span><span class="p">)</span>
    <span class="c1"># Cases of models with no input args</span>
    <span class="k">except</span> <span class="ne">IndexError</span><span class="p">:</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;No input args, skipping _decide_input_format&quot;</span><span class="p">)</span>
    <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Skipping _decide_input_format</span><span class="se">\n</span><span class="s2"> </span><span class="si">{</span><span class="n">e</span><span class="o">.</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">args</span>


<span class="k">def</span> <span class="nf">_trace</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="n">args</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">,</span> <span class="n">return_outs</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="c1"># Special case for common case of passing a single Tensor</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</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">args</span> <span class="o">=</span> <span class="p">(</span><span class="n">args</span><span class="p">,)</span>

    <span class="n">trace_graph</span><span class="p">,</span> <span class="n">torch_out</span><span class="p">,</span> <span class="n">inputs_states</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_get_trace_graph</span><span class="p">(</span>
        <span class="n">func</span><span class="p">,</span>
        <span class="n">args</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="n">_force_outplace</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">_return_inputs_states</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="n">warn_on_static_input_change</span><span class="p">(</span><span class="n">inputs_states</span><span class="p">)</span>

    <span class="n">trace_graph</span> <span class="o">=</span> <span class="n">_optimize_graph</span><span class="p">(</span><span class="n">trace_graph</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">,</span> <span class="n">params_dict</span><span class="o">=</span><span class="p">{})</span>
    <span class="k">if</span> <span class="n">return_outs</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">trace_graph</span><span class="p">,</span> <span class="n">torch_out</span>
    <span class="k">return</span> <span class="n">trace_graph</span>


<span class="k">def</span> <span class="nf">_trace_and_get_graph_from_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">):</span>
    <span class="c1"># A basic sanity check: make sure the state_dict keys are the same</span>
    <span class="c1"># before and after running the model.  Fail fast!</span>
    <span class="n">orig_state_dict_keys</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_unique_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="o">.</span><span class="n">keys</span><span class="p">()</span>

    <span class="c1"># Disable Autocast cache because it replaces kernel&#39;s weight and bias</span>
    <span class="c1"># by (undesired) constants.</span>
    <span class="c1"># No perf impact for when there are reused weights since https://github.com/pytorch/pytorch/pull/85665</span>
    <span class="n">prev_autocast_cache_enabled</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_autocast_cache_enabled</span><span class="p">()</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">set_autocast_cache_enabled</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
    <span class="n">trace_graph</span><span class="p">,</span> <span class="n">torch_out</span><span class="p">,</span> <span class="n">inputs_states</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_get_trace_graph</span><span class="p">(</span>
        <span class="n">model</span><span class="p">,</span>
        <span class="n">args</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="n">_force_outplace</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">_return_inputs_states</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">set_autocast_cache_enabled</span><span class="p">(</span><span class="n">prev_autocast_cache_enabled</span><span class="p">)</span>

    <span class="n">warn_on_static_input_change</span><span class="p">(</span><span class="n">inputs_states</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">orig_state_dict_keys</span> <span class="o">!=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_unique_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">)</span><span class="o">.</span><span class="n">keys</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
            <span class="s2">&quot;state_dict changed after running the tracer; &quot;</span>
            <span class="s2">&quot;something weird is happening in your model!&quot;</span>
        <span class="p">)</span>

    <span class="k">return</span> <span class="n">trace_graph</span><span class="p">,</span> <span class="n">torch_out</span>


<span class="k">def</span> <span class="nf">_get_param_count_list</span><span class="p">(</span><span class="n">method_graph</span><span class="p">,</span> <span class="n">args_params</span><span class="p">):</span>
    <span class="n">param_count_list</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">input_</span><span class="p">,</span> <span class="n">arg_params_</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">method_graph</span><span class="o">.</span><span class="n">inputs</span><span class="p">(),</span> <span class="n">args_params</span><span class="p">):</span>
        <span class="k">if</span> <span class="s2">&quot;PackedParams&quot;</span> <span class="ow">in</span> <span class="nb">str</span><span class="p">(</span><span class="n">input_</span><span class="o">.</span><span class="n">type</span><span class="p">()):</span>
            <span class="n">in_vars</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_flatten</span><span class="p">(</span><span class="n">arg_params_</span><span class="p">)</span>
            <span class="n">param_count_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">in_vars</span><span class="p">))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">param_count_list</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">arg_params_</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span>

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


<span class="k">def</span> <span class="nf">_check_flatten_did_not_remove</span><span class="p">(</span><span class="n">original</span><span class="p">,</span> <span class="n">jit_flattened</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;torch.jit._flatten removes None. Check if it did so in this case.&quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">flatten</span><span class="p">(</span><span class="n">x</span><span class="p">):</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
            <span class="k">for</span> <span class="n">inner</span> <span class="ow">in</span> <span class="n">x</span><span class="p">:</span>
                <span class="k">yield from</span> <span class="n">flatten</span><span class="p">(</span><span class="n">inner</span><span class="p">)</span>
        <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">inner</span> <span class="ow">in</span> <span class="n">x</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
                <span class="k">yield from</span> <span class="n">flatten</span><span class="p">(</span><span class="n">inner</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">yield</span> <span class="n">x</span>

    <span class="n">flattened_with_none</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">flatten</span><span class="p">(</span><span class="n">original</span><span class="p">))</span>
    <span class="n">num_none</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">flattened_with_none</span><span class="p">)</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</span><span class="n">jit_flattened</span><span class="p">)</span>
    <span class="k">assert</span> <span class="n">num_none</span> <span class="o">&gt;=</span> <span class="mi">0</span>
    <span class="k">if</span> <span class="n">num_none</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="sa">f</span><span class="s2">&quot;args contained </span><span class="si">{</span><span class="n">num_none</span><span class="si">}</span><span class="s2"> None&#39;s after flattening. &quot;</span>
            <span class="s2">&quot;When exporting a ScriptModule or ScriptFunction, no args may &quot;</span>
            <span class="s2">&quot;be None because that breaks type propagation.&quot;</span>
        <span class="p">)</span>


<span class="k">def</span> <span class="nf">_create_jit_graph</span><span class="p">(</span>
    <span class="n">model</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="o">|</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptFunction</span><span class="p">,</span> <span class="n">args</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Any</span><span class="p">]</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">Graph</span><span class="p">,</span> <span class="nb">list</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">IValue</span><span class="p">],</span> <span class="n">Any</span> <span class="o">|</span> <span class="kc">None</span><span class="p">,</span> <span class="n">_C</span><span class="o">.</span><span class="n">ScriptModule</span> <span class="o">|</span> <span class="kc">None</span><span class="p">]:</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptFunction</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">)):</span>
        <span class="n">flattened_args</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_flatten</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">args</span><span class="p">))[</span><span class="mi">0</span><span class="p">])</span>
        <span class="n">_check_flatten_did_not_remove</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">flattened_args</span><span class="p">)</span>
        <span class="n">torch_out</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">):</span>
            <span class="k">try</span><span class="p">:</span>
                <span class="n">graph</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">forward</span><span class="o">.</span><span class="n">graph</span>  <span class="c1"># type: ignore[attr-defined]</span>
            <span class="k">except</span> <span class="ne">AttributeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;&#39;forward&#39; method must be a script method&quot;</span><span class="p">)</span> <span class="kn">from</span> <span class="nn">e</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_function_substitution</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
            <span class="n">freezed_module</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_freeze_module</span><span class="p">(</span>
                <span class="n">cast</span><span class="p">(</span><span class="n">_C</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">,</span> <span class="n">model</span><span class="o">.</span><span class="n">_c</span><span class="p">),</span> <span class="n">preserveParameters</span><span class="o">=</span><span class="kc">True</span>
            <span class="p">)</span>
            <span class="n">module</span><span class="p">,</span> <span class="n">params</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_onnx_list_model_parameters</span><span class="p">(</span><span class="n">freezed_module</span><span class="p">)</span>
            <span class="n">method_graph</span> <span class="o">=</span> <span class="n">module</span><span class="o">.</span><span class="n">_get_method</span><span class="p">(</span><span class="s2">&quot;forward&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">graph</span>
            <span class="n">args_params</span> <span class="o">=</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">args</span><span class="p">)</span> <span class="o">+</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">params</span><span class="p">)</span>
            <span class="n">param_count_list</span> <span class="o">=</span> <span class="n">_get_param_count_list</span><span class="p">(</span><span class="n">method_graph</span><span class="p">,</span> <span class="n">args_params</span><span class="p">)</span>
            <span class="n">in_vars</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_flatten</span><span class="p">(</span><span class="n">args_params</span><span class="p">)</span>
            <span class="n">graph</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_propagate_and_assign_input_shapes</span><span class="p">(</span>
                <span class="n">method_graph</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">in_vars</span><span class="p">),</span> <span class="n">param_count_list</span><span class="p">,</span> <span class="kc">False</span><span class="p">,</span> <span class="kc">False</span>
            <span class="p">)</span>
            <span class="k">return</span> <span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">torch_out</span><span class="p">,</span> <span class="n">module</span>

        <span class="c1"># torch.jit.ScriptFunction</span>
        <span class="n">params</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">graph</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">graph</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_function_substitution</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
        <span class="n">param_count_list</span> <span class="o">=</span> <span class="n">_get_param_count_list</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
        <span class="n">graph</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_propagate_and_assign_input_shapes</span><span class="p">(</span>
            <span class="n">graph</span><span class="p">,</span> <span class="n">flattened_args</span><span class="p">,</span> <span class="n">param_count_list</span><span class="p">,</span> <span class="kc">False</span><span class="p">,</span> <span class="kc">False</span>
        <span class="p">)</span>
        <span class="k">return</span> <span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">torch_out</span><span class="p">,</span> <span class="kc">None</span>

    <span class="n">graph</span><span class="p">,</span> <span class="n">torch_out</span> <span class="o">=</span> <span class="n">_trace_and_get_graph_from_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_lint</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">state_dict</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_unique_state_dict</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
    <span class="n">params</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">state_dict</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
    <span class="n">graph_inputs</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">graph</span><span class="o">.</span><span class="n">inputs</span><span class="p">())</span>
    <span class="n">user_input_num</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">graph_inputs</span><span class="p">)</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</span><span class="n">state_dict</span><span class="p">)</span>
    <span class="n">param_names</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">state_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
    <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">inp</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">graph_inputs</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">i</span> <span class="o">&gt;=</span> <span class="n">user_input_num</span><span class="p">:</span>
            <span class="n">inp</span><span class="o">.</span><span class="n">setDebugName</span><span class="p">(</span><span class="n">param_names</span><span class="p">[</span><span class="n">i</span> <span class="o">-</span> <span class="n">user_input_num</span><span class="p">])</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_function_substitution</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">torch_out</span><span class="p">,</span> <span class="kc">None</span>


<span class="k">def</span> <span class="nf">_get_named_param_dict</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">):</span>
    <span class="n">input_and_param_names</span> <span class="o">=</span> <span class="p">[</span><span class="n">val</span><span class="o">.</span><span class="n">debugName</span><span class="p">()</span> <span class="k">for</span> <span class="n">val</span> <span class="ow">in</span> <span class="n">graph</span><span class="o">.</span><span class="n">inputs</span><span class="p">()]</span>
    <span class="n">param_names</span> <span class="o">=</span> <span class="n">input_and_param_names</span><span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">input_and_param_names</span><span class="p">)</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="p">:]</span>
    <span class="n">_params_dict</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">param_names</span><span class="p">,</span> <span class="n">params</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_params_dict</span>


<span class="k">def</span> <span class="nf">_get_example_outputs</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">):</span>
    <span class="n">input_args</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">args</span><span class="p">)</span>
    <span class="n">input_kwargs</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">if</span> <span class="n">input_args</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">input_args</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="nb">dict</span><span class="p">):</span>
        <span class="n">input_kwargs</span> <span class="o">=</span> <span class="n">input_args</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">input_args</span> <span class="o">=</span> <span class="n">input_args</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>

    <span class="n">example_outputs</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="o">*</span><span class="n">input_args</span><span class="p">,</span> <span class="o">**</span><span class="n">input_kwargs</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">example_outputs</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
        <span class="n">example_outputs</span> <span class="o">=</span> <span class="p">[</span><span class="n">example_outputs</span><span class="p">]</span>
    <span class="k">elif</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">example_outputs</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
        <span class="n">example_outputs</span> <span class="o">=</span> <span class="p">(</span><span class="n">example_outputs</span><span class="p">,)</span>

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


<span class="n">_qtype_vtype_map</span> <span class="o">=</span> <span class="p">{</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">quint8</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">uint8</span><span class="p">,</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">qint8</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">qint32</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">int32</span><span class="p">,</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">quint4x2</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">int8</span><span class="p">,</span>
<span class="p">}</span>


<span class="k">def</span> <span class="nf">unpack_quantized_tensor</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">cast_onnx_accepted</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</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="ow">and</span> <span class="n">value</span><span class="o">.</span><span class="n">dtype</span> <span class="ow">in</span> <span class="n">_qtype_vtype_map</span><span class="p">:</span>
        <span class="n">q_value_dequantize</span> <span class="o">=</span> <span class="n">value</span><span class="o">.</span><span class="n">dequantize</span><span class="p">()</span>
        <span class="n">q_scale</span> <span class="o">=</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">value</span><span class="o">.</span><span class="n">q_scale</span><span class="p">(),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">double</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">cast_onnx_accepted</span>
            <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">value</span><span class="o">.</span><span class="n">q_scale</span><span class="p">(),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
        <span class="p">)</span>
        <span class="n">q_zero_point</span> <span class="o">=</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">value</span><span class="o">.</span><span class="n">q_zero_point</span><span class="p">(),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">int64</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">cast_onnx_accepted</span>
            <span class="k">else</span> <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">value</span><span class="o">.</span><span class="n">q_zero_point</span><span class="p">(),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">_qtype_vtype_map</span><span class="p">[</span><span class="n">value</span><span class="o">.</span><span class="n">dtype</span><span class="p">])</span>
        <span class="p">)</span>
        <span class="n">q_value</span> <span class="o">=</span> <span class="n">q_value_dequantize</span> <span class="o">/</span> <span class="n">q_scale</span> <span class="o">+</span> <span class="n">q_zero_point</span>
        <span class="n">q_value</span> <span class="o">=</span> <span class="n">q_value</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="n">dtype</span><span class="o">=</span><span class="n">_qtype_vtype_map</span><span class="p">[</span><span class="n">value</span><span class="o">.</span><span class="n">dtype</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">q_value</span><span class="p">,</span> <span class="n">q_scale</span><span class="p">,</span> <span class="n">q_zero_point</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="p">(</span><span class="n">value</span><span class="p">,)</span>


<span class="k">def</span> <span class="nf">_pre_trace_quant_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Returns `torch.jit.trace(model, args)` if model is quantized. Otherwise do nothing and return</span>
<span class="sd">    original model.</span>

<span class="sd">    This is due to https://github.com/pytorch/pytorch/issues/75761.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">any</span><span class="p">(</span>
        <span class="nb">hasattr</span><span class="p">(</span><span class="n">m</span><span class="p">,</span> <span class="s2">&quot;_packed_params&quot;</span><span class="p">)</span> <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;modules&quot;</span><span class="p">,</span> <span class="nb">list</span><span class="p">)()</span>
    <span class="p">)</span> <span class="ow">or</span> <span class="nb">any</span><span class="p">(</span><span class="nb">getattr</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="s2">&quot;is_quantized&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span> <span class="k">for</span> <span class="n">arg</span> <span class="ow">in</span> <span class="n">args</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">model</span>


<span class="k">def</span> <span class="nf">_model_to_graph</span><span class="p">(</span>
    <span class="n">model</span><span class="p">,</span>
    <span class="n">args</span><span class="p">,</span>
    <span class="n">verbose</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">input_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">output_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">operator_export_type</span><span class="o">=</span><span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span><span class="p">,</span>
    <span class="n">do_constant_folding</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="n">_disable_torch_constant_prop</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">fixed_batch_size</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">training</span><span class="o">=</span><span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span><span class="p">,</span>
    <span class="n">dynamic_axes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">Graph</span><span class="p">,</span>
    <span class="nb">dict</span><span class="p">[</span><span class="nb">str</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">torch</span><span class="o">.</span><span class="n">Tensor</span>
    <span class="o">|</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span>
    <span class="o">|</span> <span class="nb">list</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="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span>
    <span class="o">|</span> <span class="n">Any</span>
    <span class="o">|</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">]:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Converts model into an ONNX graph.</span>

<span class="sd">    Returns:</span>
<span class="sd">        graph: A TorchScript IR Graph with ONNX nodes.</span>
<span class="sd">        params_dict: Dict from input param name to param value.</span>
<span class="sd">        torch_out: The output tensors resulting from the trace of ``model``.</span>
<span class="sd">            If ``model`` is a :class:`torch.jit.ScriptModule` or :class:`torch.jit.ScriptFunction`,</span>
<span class="sd">            this will be None, since we are not doing any tracing.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># TODO: can we simplify this to always return a tuple of Tensor or None?</span>

    <span class="c1"># Special case for common case of passing a single Tensor</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">args</span><span class="p">,</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="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="nb">bool</span><span class="p">)):</span>
        <span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="n">args</span><span class="p">,)</span>

    <span class="n">model</span> <span class="o">=</span> <span class="n">_pre_trace_quant_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
    <span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">,</span> <span class="n">torch_out</span><span class="p">,</span> <span class="n">module</span> <span class="o">=</span> <span class="n">_create_jit_graph</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
    <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_get_named_param_dict</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span>

    <span class="k">try</span><span class="p">:</span>
        <span class="n">graph</span> <span class="o">=</span> <span class="n">_optimize_graph</span><span class="p">(</span>
            <span class="n">graph</span><span class="p">,</span>
            <span class="n">operator_export_type</span><span class="p">,</span>
            <span class="n">_disable_torch_constant_prop</span><span class="o">=</span><span class="n">_disable_torch_constant_prop</span><span class="p">,</span>
            <span class="n">fixed_batch_size</span><span class="o">=</span><span class="n">fixed_batch_size</span><span class="p">,</span>
            <span class="n">params_dict</span><span class="o">=</span><span class="n">params_dict</span><span class="p">,</span>
            <span class="n">dynamic_axes</span><span class="o">=</span><span class="n">dynamic_axes</span><span class="p">,</span>
            <span class="n">input_names</span><span class="o">=</span><span class="n">input_names</span><span class="p">,</span>
            <span class="n">module</span><span class="o">=</span><span class="n">module</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_onnx_log</span><span class="p">(</span><span class="s2">&quot;Torch IR graph at exception: &quot;</span><span class="p">,</span> <span class="n">graph</span><span class="p">)</span>
        <span class="k">raise</span>

    <span class="n">is_script</span> <span class="o">=</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptFunction</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">))</span>
    <span class="k">if</span> <span class="n">is_script</span><span class="p">:</span>
        <span class="n">example_outputs</span> <span class="o">=</span> <span class="n">_get_example_outputs</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
        <span class="n">example_outputs_final</span> <span class="o">=</span> <span class="p">()</span>
        <span class="k">for</span> <span class="n">example_output</span> <span class="ow">in</span> <span class="n">example_outputs</span><span class="p">:</span>
            <span class="n">example_outputs_final</span> <span class="o">+=</span> <span class="n">unpack_quantized_tensor</span><span class="p">(</span><span class="n">example_output</span><span class="p">)</span>
        <span class="n">out_vars</span><span class="p">,</span> <span class="n">desc</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_flatten</span><span class="p">(</span><span class="n">example_outputs_final</span><span class="p">)</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_assign_output_shape</span><span class="p">(</span>
            <span class="n">graph</span><span class="p">,</span>
            <span class="n">out_vars</span><span class="p">,</span>
            <span class="n">desc</span><span class="p">,</span>
            <span class="n">GLOBALS</span><span class="o">.</span><span class="n">onnx_shape_inference</span><span class="p">,</span>
            <span class="n">is_script</span><span class="p">,</span>
            <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="c1"># NB: ONNX requires complete information about output types, which might be</span>
    <span class="c1"># erased by some optimizations, so we need to set it explicitly again.</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">torch_out</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
            <span class="n">output_wrapped</span> <span class="o">=</span> <span class="p">[</span><span class="n">torch_out</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">output_wrapped</span> <span class="o">=</span> <span class="n">torch_out</span>  <span class="c1"># type: ignore[assignment]</span>

        <span class="n">output_tensors</span><span class="p">,</span> <span class="n">out_desc</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_flatten</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">output_wrapped</span><span class="p">))</span>
        <span class="c1"># assign_output_shape pass is not compatible with quantized outputs.</span>
        <span class="c1"># Quantized outputs are flattened to 3 values in ONNX, while packed as</span>
        <span class="c1"># single value in PyTorch.</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="nb">getattr</span><span class="p">(</span><span class="n">out</span><span class="p">,</span> <span class="s2">&quot;is_quantized&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span> <span class="k">for</span> <span class="n">out</span> <span class="ow">in</span> <span class="n">output_tensors</span><span class="p">):</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_assign_output_shape</span><span class="p">(</span>
                <span class="n">graph</span><span class="p">,</span>
                <span class="n">output_tensors</span><span class="p">,</span>
                <span class="n">out_desc</span><span class="p">,</span>
                <span class="n">GLOBALS</span><span class="o">.</span><span class="n">onnx_shape_inference</span><span class="p">,</span>
                <span class="n">is_script</span><span class="p">,</span>
                <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span><span class="p">,</span>
            <span class="p">)</span>

    <span class="n">_set_input_and_output_names</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">input_names</span><span class="p">,</span> <span class="n">output_names</span><span class="p">)</span>
    <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_get_named_param_dict</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">)</span>

    <span class="k">if</span> <span class="p">(</span>
        <span class="n">do_constant_folding</span>
        <span class="ow">and</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span>
        <span class="o">&gt;=</span> <span class="n">_constants</span><span class="o">.</span><span class="n">ONNX_CONSTANT_FOLDING_MIN_OPSET</span>
    <span class="p">):</span>
        <span class="k">if</span> <span class="n">training</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">training</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span><span class="p">:</span>
            <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_eval_peephole</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">)</span>

        <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_constant_fold</span><span class="p">(</span>
            <span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">,</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span>
        <span class="p">)</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_dce_allow_deleting_nodes_with_side_effects</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">onnx_shape_inference</span><span class="p">:</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_graph_shape_type_inference</span><span class="p">(</span>
                <span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">,</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span>
            <span class="p">)</span>
        <span class="k">except</span> <span class="ne">RuntimeError</span><span class="p">:</span>
            <span class="c1"># NOTE: shape type inference error should not stop the export process</span>
            <span class="c1"># https://github.com/pytorch/pytorch/issues/132205</span>
            <span class="k">pass</span>

    <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_eliminate_unused_items</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">)</span>

    <span class="c1"># For ONNX opset &lt; 9, constants only have three data types: float16, float, double.</span>
    <span class="c1"># In this pass transform constants of other data types to float/double + cast operator.</span>
    <span class="k">if</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span> <span class="o">&lt;</span> <span class="mi">9</span><span class="p">:</span>
        <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_cast_all_constant_to_floating</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_filter_non_tensor_arguments</span><span class="p">(</span><span class="n">params_dict</span><span class="p">)</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_decay_packed_param_input_types</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="c1"># If output names lack a proper name and are identified only by their unique</span>
    <span class="c1"># give them a legible name for debugging purposes</span>
    <span class="n">_apply_friendly_debug_names</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">,</span> <span class="n">torch_out</span>


<span class="nd">@_deprecation</span><span class="o">.</span><span class="n">deprecated</span><span class="p">(</span><span class="s2">&quot;2.5&quot;</span><span class="p">,</span> <span class="s2">&quot;the future&quot;</span><span class="p">,</span> <span class="s2">&quot;avoid using this function&quot;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">unconvertible_ops</span><span class="p">(</span>
    <span class="n">model</span><span class="p">,</span>
    <span class="n">args</span><span class="p">,</span>
    <span class="n">training</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span> <span class="o">=</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span><span class="p">,</span>
    <span class="n">opset_version</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">tuple</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">Graph</span><span class="p">,</span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">]]:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Returns an approximated list of all ops that are yet supported by :mod:`torch.onnx`.</span>

<span class="sd">    The list is approximated because some ops may be removed during the conversion</span>
<span class="sd">    process and don&#39;t need to be converted. Some other ops may have partial support</span>
<span class="sd">    that will fail conversion with particular inputs. Please open a Github Issue</span>
<span class="sd">    for op support requests.</span>

<span class="sd">    Args:</span>
<span class="sd">        model: Same as the `model` parameter in :func:`torch.onnx.export`.</span>
<span class="sd">        args: Same as the `args` parameter in :func:`torch.onnx.export`.</span>
<span class="sd">        training: Same as the `training` parameter in :func:`torch.onnx.export`.</span>
<span class="sd">        opset_version: Same as the `opset_version` parameter in :func:`torch.onnx.export`.</span>

<span class="sd">    Returns:</span>
<span class="sd">        The JIT graph and a list of unconvertible ops in the format of &quot;domain::op&quot;.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">opset_version</span> <span class="o">=</span> <span class="n">opset_version</span> <span class="ow">or</span> <span class="n">_constants</span><span class="o">.</span><span class="n">ONNX_DEFAULT_OPSET</span>
    <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span> <span class="o">=</span> <span class="n">opset_version</span>

    <span class="k">try</span><span class="p">:</span>
        <span class="k">with</span> <span class="n">exporter_context</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">training</span><span class="p">,</span> <span class="n">verbose</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
            <span class="c1"># Create a mostly clean JIT graph that contains the plain aten and</span>
            <span class="c1"># other ops we can check with the symbolic registry.</span>
            <span class="c1"># NOTE: We don&#39;t want to actually convert any ops to ONNX or run any</span>
            <span class="c1"># symbolic functions because there is a higher chance that a pass</span>
            <span class="c1"># fails or an unconvertible op messes up the graph during ONNX conversion.</span>
            <span class="c1"># This way we can always generate a list just by looking at the names</span>
            <span class="c1"># of the ops in the graph.</span>
            <span class="n">args</span> <span class="o">=</span> <span class="n">_decide_input_format</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
            <span class="n">model</span> <span class="o">=</span> <span class="n">_pre_trace_quant_model</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
            <span class="n">graph</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">module</span> <span class="o">=</span> <span class="n">_create_jit_graph</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_inline</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_remove_inplace_ops_for_onnx</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">module</span><span class="p">)</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_erase_number_types</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_dce_allow_deleting_nodes_with_side_effects</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="k">raise</span> <span class="n">errors</span><span class="o">.</span><span class="n">OnnxExporterError</span><span class="p">(</span>
            <span class="s2">&quot;Failed to discover unconvertible ops because of errors during the JIT graph &quot;</span>
            <span class="s2">&quot;generation process.&quot;</span>
        <span class="p">)</span> <span class="kn">from</span> <span class="nn">e</span>

    <span class="n">unsupported_ops</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">graph</span><span class="o">.</span><span class="n">nodes</span><span class="p">():</span>
        <span class="n">domain_op</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">kind</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">domain_op</span><span class="o">.</span><span class="n">startswith</span><span class="p">((</span><span class="s2">&quot;onnx::&quot;</span><span class="p">,</span> <span class="s2">&quot;prim::&quot;</span><span class="p">)):</span>
            <span class="c1"># We consider onnx and prim ops as supported ops, even though some &quot;prim&quot;</span>
            <span class="c1"># ops are not implemented as symbolic functions, because they may be</span>
            <span class="c1"># eliminated in the conversion passes. Users may still see errors caused</span>
            <span class="c1"># by prim ops even though they don&#39;t show up in the list.</span>
            <span class="k">continue</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">registration</span><span class="o">.</span><span class="n">registry</span><span class="o">.</span><span class="n">is_registered_op</span><span class="p">(</span>
            <span class="n">domain_op</span><span class="o">.</span><span class="n">rstrip</span><span class="p">(</span><span class="s2">&quot;_&quot;</span><span class="p">),</span> <span class="n">opset_version</span>
        <span class="p">):</span>
            <span class="c1"># We consider all registered ops supported, even though some of them are</span>
            <span class="c1"># only partially supported, because there is not yet a good way to check</span>
            <span class="c1"># if an op is fully supported.</span>
            <span class="c1"># TODO(justinchuby): Create a way to check if an op is fully supported.</span>
            <span class="n">unsupported_ops</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">domain_op</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">graph</span><span class="p">,</span> <span class="n">unsupported_ops</span>


<span class="k">def</span> <span class="nf">_setup_trace_module_map</span><span class="p">(</span>
    <span class="n">model</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="o">|</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">,</span>
    <span class="n">export_modules_as_functions</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">|</span> <span class="n">Collection</span><span class="p">[</span><span class="nb">type</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="p">)</span> <span class="o">-&gt;</span> <span class="nb">set</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span>
    <span class="k">def</span> <span class="nf">__register_attribute_hook</span><span class="p">():</span>
        <span class="n">attr_name</span> <span class="o">=</span> <span class="s2">&quot;_onnx_attrs&quot;</span>

        <span class="k">def</span> <span class="nf">_track_module_attributes_forward_pre_hook</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="nb">input</span><span class="p">):</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">attr_name</span><span class="p">,</span> <span class="n">_get_module_attributes</span><span class="p">(</span><span class="n">module</span><span class="p">))</span>

        <span class="k">def</span> <span class="nf">_track_module_attributes_forward_hook</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output</span><span class="p">):</span>
            <span class="n">tracing_state</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_get_tracing_state</span><span class="p">()</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">tracing_state</span><span class="p">:</span>
                <span class="k">return</span>

            <span class="n">graph</span> <span class="o">=</span> <span class="n">tracing_state</span><span class="o">.</span><span class="n">graph</span><span class="p">()</span>
            <span class="n">onnx_attrs</span> <span class="o">=</span> <span class="p">{}</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">attr_name</span><span class="p">):</span>
                <span class="n">onnx_attrs</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">attr_name</span><span class="p">)</span>
                <span class="nb">delattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">attr_name</span><span class="p">)</span>

            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_track_scope_attributes</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">onnx_attrs</span><span class="p">)</span>

        <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">modules</span><span class="p">():</span>
            <span class="n">m</span><span class="o">.</span><span class="n">register_forward_hook</span><span class="p">(</span><span class="n">_track_module_attributes_forward_hook</span><span class="p">)</span>
            <span class="n">m</span><span class="o">.</span><span class="n">register_forward_pre_hook</span><span class="p">(</span><span class="n">_track_module_attributes_forward_pre_hook</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">_unqualified_variable_name</span><span class="p">(</span><span class="n">qualified_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Parse qualified variable name and return the unqualified version.</span>

<span class="sd">        Pure numeric atoms are considered inadequate, so this function will look past them,</span>
<span class="sd">        and start from the first non-numeric atom.</span>

<span class="sd">        Example:</span>
<span class="sd">            &gt;&gt;&gt; _unqualified_variable_name(&quot;__main__.Foo.bar&quot;)</span>
<span class="sd">            &#39;bar&#39;</span>
<span class="sd">            &gt;&gt;&gt; _unqualified_variable_name(&quot;__main__.Foo.bar.0&quot;)</span>
<span class="sd">            &#39;bar.0&#39;</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">name_atoms</span> <span class="o">=</span> <span class="n">qualified_name</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s2">&quot;.&quot;</span><span class="p">)</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">atom</span> <span class="ow">in</span> <span class="nb">reversed</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="nb">enumerate</span><span class="p">(</span><span class="n">name_atoms</span><span class="p">))):</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">atom</span><span class="o">.</span><span class="n">isnumeric</span><span class="p">():</span>
                <span class="k">return</span> <span class="s2">&quot;.&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">name_atoms</span><span class="p">[</span><span class="n">i</span><span class="p">:])</span>
        <span class="k">return</span> <span class="n">qualified_name</span>

    <span class="n">trace_module_map</span> <span class="o">=</span> <span class="p">{</span>
        <span class="n">_m</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_jit_onnx_create_full_scope_name</span><span class="p">(</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">typename</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">_m</span><span class="p">)),</span> <span class="n">_unqualified_variable_name</span><span class="p">(</span><span class="n">_n</span><span class="p">)</span>
        <span class="p">)</span>
        <span class="k">for</span> <span class="n">_n</span><span class="p">,</span> <span class="n">_m</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">named_modules</span><span class="p">()</span>
    <span class="p">}</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_trace</span><span class="o">.</span><span class="n">_trace_module_map</span> <span class="o">=</span> <span class="n">trace_module_map</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">export_modules_as_functions</span><span class="p">,</span> <span class="nb">bool</span><span class="p">)</span> <span class="ow">and</span> <span class="n">export_modules_as_functions</span><span class="p">:</span>
        <span class="n">module_typenames</span> <span class="o">=</span> <span class="p">{</span><span class="n">torch</span><span class="o">.</span><span class="n">typename</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">module</span><span class="p">))</span> <span class="k">for</span> <span class="n">module</span> <span class="ow">in</span> <span class="n">trace_module_map</span><span class="p">}</span>
    <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">export_modules_as_functions</span><span class="p">,</span> <span class="nb">set</span><span class="p">)</span> <span class="ow">and</span> <span class="n">export_modules_as_functions</span><span class="p">:</span>

        <span class="k">def</span> <span class="nf">_find_typename</span><span class="p">(</span><span class="n">v</span><span class="p">):</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">v</span><span class="p">,</span> <span class="nb">type</span><span class="p">):</span>
                <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">typename</span><span class="p">(</span><span class="n">v</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                    <span class="s2">&quot;Only type of the `nn.Module` should be &quot;</span>
                    <span class="s2">&quot;passed in the set for argument `export_modules_as_functions`. &quot;</span>
                    <span class="sa">f</span><span class="s2">&quot;Got `</span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">v</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">`.&quot;</span>
                <span class="p">)</span>

        <span class="n">module_typenames</span> <span class="o">=</span> <span class="p">{</span><span class="n">_find_typename</span><span class="p">(</span><span class="n">v</span><span class="p">)</span> <span class="k">for</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">export_modules_as_functions</span><span class="p">}</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">module_typenames</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">module_typenames</span><span class="p">:</span>
        <span class="n">__register_attribute_hook</span><span class="p">()</span>

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


<span class="k">def</span> <span class="nf">_reset_trace_module_map</span><span class="p">():</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">_trace</span><span class="o">.</span><span class="n">_trace_module_map</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_clear_scope_records</span><span class="p">()</span>


<span class="k">def</span> <span class="nf">_get_module_attributes</span><span class="p">(</span><span class="n">module</span><span class="p">):</span>
    <span class="n">annotations</span> <span class="o">=</span> <span class="n">typing</span><span class="o">.</span><span class="n">get_type_hints</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">module</span><span class="p">))</span>
    <span class="n">base_m_annotations</span> <span class="o">=</span> <span class="n">typing</span><span class="o">.</span><span class="n">get_type_hints</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="p">[</span><span class="n">annotations</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">base_m_annotations</span><span class="p">]</span>
    <span class="c1"># Check whether module attributes can be accessed. Some classes</span>
    <span class="c1"># define attributes but don&#39;t provide access to them in their</span>
    <span class="c1"># constructor.</span>
    <span class="c1">#</span>
    <span class="c1"># For example, torch.nn.Embedding has the `freeze` variable and its</span>
    <span class="c1"># type specified in the class but the attribute is not created in the</span>
    <span class="c1"># constructor. In other words, there is no `self.freeze = &lt;True | False&gt;`</span>
    <span class="c1"># in the constructor.</span>
    <span class="c1">#</span>
    <span class="c1"># Reference: https://github.com/pytorch/pytorch/blob/92de1d322223fb5584e384971b32c46b93bc2f4b/torch/nn/modules/sparse.py#L120</span>
    <span class="n">attrs</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">annotations</span><span class="p">:</span>
        <span class="k">try</span><span class="p">:</span>
            <span class="n">attrs</span><span class="p">[</span><span class="n">k</span><span class="p">]</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">module</span><span class="p">,</span> <span class="n">k</span><span class="p">)</span>
        <span class="k">except</span> <span class="ne">AttributeError</span><span class="p">:</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_onnx_log</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Skipping module attribute &#39;</span><span class="si">{</span><span class="n">k</span><span class="si">}</span><span class="s2">&#39;&quot;</span><span class="p">)</span>
            <span class="k">continue</span>
    <span class="k">return</span> <span class="n">attrs</span>


<span class="k">def</span> <span class="nf">_export</span><span class="p">(</span>
    <span class="n">model</span><span class="p">,</span>
    <span class="n">args</span><span class="p">,</span>
    <span class="n">f</span><span class="p">,</span>
    <span class="n">export_params</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="n">verbose</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">training</span><span class="o">=</span><span class="n">_C_onnx</span><span class="o">.</span><span class="n">TrainingMode</span><span class="o">.</span><span class="n">EVAL</span><span class="p">,</span>
    <span class="n">input_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">output_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">operator_export_type</span><span class="o">=</span><span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span><span class="p">,</span>
    <span class="n">export_type</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">opset_version</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">do_constant_folding</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="n">dynamic_axes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">keep_initializers_as_inputs</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">fixed_batch_size</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">custom_opsets</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">add_node_names</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="n">onnx_shape_inference</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="n">export_modules_as_functions</span><span class="p">:</span> <span class="n">Any</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
    <span class="n">autograd_inlining</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
<span class="p">):</span>
    <span class="k">assert</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">in_onnx_export</span> <span class="ow">is</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</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">DataParallel</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;torch.nn.DataParallel is not supported by ONNX &quot;</span>
            <span class="s2">&quot;exporter, please use &#39;attribute&#39; module to &quot;</span>
            <span class="s2">&quot;unwrap model from torch.nn.DataParallel. Try &quot;</span>
            <span class="s2">&quot;torch.onnx.export(model.module, ...)&quot;</span>
        <span class="p">)</span>

    <span class="n">GLOBALS</span><span class="o">.</span><span class="n">onnx_shape_inference</span> <span class="o">=</span> <span class="n">onnx_shape_inference</span>

    <span class="k">if</span> <span class="n">opset_version</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">opset_version</span> <span class="o">=</span> <span class="n">_constants</span><span class="o">.</span><span class="n">ONNX_DEFAULT_OPSET</span>

    <span class="c1"># torch.onnx.export does not support opset versions &gt;=18</span>
    <span class="k">if</span> <span class="n">opset_version</span> <span class="o">&gt;</span> <span class="n">_constants</span><span class="o">.</span><span class="n">ONNX_TORCHSCRIPT_EXPORTER_MAX_OPSET</span><span class="p">:</span>
        <span class="c1"># We do not want to fail because we should still allow users to create</span>
        <span class="c1"># custom symbolic functions for opset&gt;17</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
            <span class="sa">f</span><span class="s2">&quot;Exporting to ONNX opset version </span><span class="si">{</span><span class="n">opset_version</span><span class="si">}</span><span class="s2"> is not supported. &quot;</span>
            <span class="sa">f</span><span class="s2">&quot;by &#39;torch.onnx.export()&#39;. &quot;</span>
            <span class="sa">f</span><span class="s2">&quot;The highest opset version supported is </span><span class="si">{</span><span class="n">_constants</span><span class="o">.</span><span class="n">ONNX_TORCHSCRIPT_EXPORTER_MAX_OPSET</span><span class="si">}</span><span class="s2">. &quot;</span>
            <span class="sa">f</span><span class="s2">&quot;To use a newer opset version, consider &#39;torch.onnx.export(..., dynamo=True)&#39;. &quot;</span><span class="p">,</span>
            <span class="n">category</span><span class="o">=</span><span class="n">errors</span><span class="o">.</span><span class="n">OnnxExporterWarning</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">export_modules_as_functions</span> <span class="ow">and</span> <span class="n">opset_version</span> <span class="o">&lt;</span> <span class="mi">15</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;`export_modules_as_functions` is not supported for `opset_version` &lt; 15.&quot;</span>
            <span class="s2">&quot;This is because `opset_version` &lt; 15 implies IR version &lt; 8, which means &quot;</span>
            <span class="s2">&quot;no local function support. &quot;</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">operator_export_type</span><span class="p">:</span>
        <span class="n">operator_export_type</span> <span class="o">=</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span>

    <span class="c1"># By default, training=TrainingMode.EVAL,</span>
    <span class="c1"># which is good because running a model in training mode could result in</span>
    <span class="c1"># internal buffers getting updated, dropout getting applied, etc.</span>
    <span class="c1"># If you really know what you&#39;re doing, you can turn</span>
    <span class="c1"># training=TrainingMode.TRAINING or training=TrainingMode.PRESERVE,</span>
    <span class="c1"># (to preserve whatever the original training mode was.)</span>
    <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span> <span class="o">=</span> <span class="n">opset_version</span>
    <span class="n">GLOBALS</span><span class="o">.</span><span class="n">operator_export_type</span> <span class="o">=</span> <span class="n">operator_export_type</span>

    <span class="k">try</span><span class="p">:</span>
        <span class="n">GLOBALS</span><span class="o">.</span><span class="n">in_onnx_export</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="n">_autograd_inlining_previous</span> <span class="o">=</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">autograd_inlining</span>
        <span class="n">GLOBALS</span><span class="o">.</span><span class="n">autograd_inlining</span> <span class="o">=</span> <span class="n">autograd_inlining</span>

        <span class="n">module_typenames_to_export_as_functions</span><span class="p">:</span> <span class="nb">set</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</span><span class="p">,</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="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">ScriptModule</span><span class="p">)):</span>
            <span class="n">module_typenames_to_export_as_functions</span> <span class="o">=</span> <span class="n">_setup_trace_module_map</span><span class="p">(</span>
                <span class="n">model</span><span class="p">,</span> <span class="n">export_modules_as_functions</span>
            <span class="p">)</span>

        <span class="k">with</span> <span class="n">exporter_context</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">training</span><span class="p">,</span> <span class="n">verbose</span><span class="p">):</span>
            <span class="n">val_keep_init_as_ip</span> <span class="o">=</span> <span class="n">_decide_keep_init_as_input</span><span class="p">(</span>
                <span class="n">keep_initializers_as_inputs</span><span class="p">,</span>
                <span class="n">operator_export_type</span><span class="p">,</span>
                <span class="n">opset_version</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="n">val_add_node_names</span> <span class="o">=</span> <span class="n">_decide_add_node_names</span><span class="p">(</span>
                <span class="n">add_node_names</span><span class="p">,</span> <span class="n">operator_export_type</span>
            <span class="p">)</span>
            <span class="n">val_do_constant_folding</span> <span class="o">=</span> <span class="n">_decide_constant_folding</span><span class="p">(</span>
                <span class="n">do_constant_folding</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">,</span> <span class="n">training</span>
            <span class="p">)</span>
            <span class="c1"># Normally f can be a file-like object, but for large models, the external data format requires a</span>
            <span class="c1"># valid `model_file_location`. Code in export.cpp will enforce this.</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
                <span class="n">model_file_location</span> <span class="o">=</span> <span class="n">f</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">model_file_location</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
            <span class="n">args</span> <span class="o">=</span> <span class="n">_decide_input_format</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">args</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">dynamic_axes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dynamic_axes</span> <span class="o">=</span> <span class="p">{}</span>
            <span class="n">_validate_dynamic_axes</span><span class="p">(</span><span class="n">dynamic_axes</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">input_names</span><span class="p">,</span> <span class="n">output_names</span><span class="p">)</span>

            <span class="n">graph</span><span class="p">,</span> <span class="n">params_dict</span><span class="p">,</span> <span class="n">torch_out</span> <span class="o">=</span> <span class="n">_model_to_graph</span><span class="p">(</span>
                <span class="n">model</span><span class="p">,</span>
                <span class="n">args</span><span class="p">,</span>
                <span class="n">verbose</span><span class="p">,</span>
                <span class="n">input_names</span><span class="p">,</span>
                <span class="n">output_names</span><span class="p">,</span>
                <span class="n">operator_export_type</span><span class="p">,</span>
                <span class="n">val_do_constant_folding</span><span class="p">,</span>
                <span class="n">fixed_batch_size</span><span class="o">=</span><span class="n">fixed_batch_size</span><span class="p">,</span>
                <span class="n">training</span><span class="o">=</span><span class="n">training</span><span class="p">,</span>
                <span class="n">dynamic_axes</span><span class="o">=</span><span class="n">dynamic_axes</span><span class="p">,</span>
            <span class="p">)</span>

            <span class="k">if</span> <span class="n">custom_opsets</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">custom_opsets</span> <span class="o">=</span> <span class="p">{}</span>

            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_dce_allow_deleting_nodes_with_side_effects</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
            <span class="n">node_attr_to_name</span> <span class="o">=</span> <span class="p">{}</span>  <span class="c1"># type: ignore[var-annotated]</span>
            <span class="k">if</span> <span class="n">module_typenames_to_export_as_functions</span><span class="p">:</span>
                <span class="c1"># NOTE: cannot call DCE after this pass. DCE will remove function definition nodes.</span>
                <span class="n">node_attr_to_name</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_function_extraction</span><span class="p">(</span>
                    <span class="n">graph</span><span class="p">,</span>
                    <span class="n">module_typenames_to_export_as_functions</span><span class="p">,</span>
                    <span class="nb">list</span><span class="p">(</span><span class="n">params_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">()),</span>
                <span class="p">)</span>

            <span class="k">if</span> <span class="n">keep_initializers_as_inputs</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">True</span><span class="p">:</span>
                <span class="n">params_dict</span> <span class="o">=</span> <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_deduplicate_initializers</span><span class="p">(</span>  <span class="c1"># type: ignore[assignment]</span>
                    <span class="n">graph</span><span class="p">,</span>
                    <span class="n">params_dict</span><span class="p">,</span>  <span class="c1"># type: ignore[arg-type]</span>
                    <span class="nb">getattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;training&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">),</span>  <span class="c1"># type: ignore[arg-type]</span>
                <span class="p">)</span>
            <span class="n">_C</span><span class="o">.</span><span class="n">_jit_pass_onnx_assign_scoped_names_for_node_and_value</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
            <span class="n">defer_weight_export</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="k">if</span> <span class="n">export_params</span><span class="p">:</span>
                <span class="p">(</span>
                    <span class="n">proto</span><span class="p">,</span>
                    <span class="n">export_map</span><span class="p">,</span>
                    <span class="n">_val_use_external_data_format</span><span class="p">,</span>
                    <span class="n">_node_names</span><span class="p">,</span>
                <span class="p">)</span> <span class="o">=</span> <span class="n">graph</span><span class="o">.</span><span class="n">_export_onnx</span><span class="p">(</span>  <span class="c1"># type: ignore[attr-defined]</span>
                    <span class="n">params_dict</span><span class="p">,</span>
                    <span class="n">opset_version</span><span class="p">,</span>
                    <span class="n">dynamic_axes</span><span class="p">,</span>
                    <span class="n">defer_weight_export</span><span class="p">,</span>
                    <span class="n">operator_export_type</span><span class="p">,</span>
                    <span class="ow">not</span> <span class="n">verbose</span><span class="p">,</span>
                    <span class="n">val_keep_init_as_ip</span><span class="p">,</span>
                    <span class="n">custom_opsets</span><span class="p">,</span>
                    <span class="n">val_add_node_names</span><span class="p">,</span>
                    <span class="n">model_file_location</span><span class="p">,</span>
                    <span class="n">node_attr_to_name</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="p">(</span>
                    <span class="n">proto</span><span class="p">,</span>
                    <span class="n">export_map</span><span class="p">,</span>
                    <span class="n">_</span><span class="p">,</span>
                    <span class="n">_</span><span class="p">,</span>
                <span class="p">)</span> <span class="o">=</span> <span class="n">graph</span><span class="o">.</span><span class="n">_export_onnx</span><span class="p">(</span>  <span class="c1"># type: ignore[attr-defined]</span>
                    <span class="p">{},</span>
                    <span class="n">opset_version</span><span class="p">,</span>
                    <span class="n">dynamic_axes</span><span class="p">,</span>
                    <span class="n">defer_weight_export</span><span class="p">,</span>
                    <span class="n">operator_export_type</span><span class="p">,</span>
                    <span class="ow">not</span> <span class="n">verbose</span><span class="p">,</span>
                    <span class="n">val_keep_init_as_ip</span><span class="p">,</span>
                    <span class="n">custom_opsets</span><span class="p">,</span>
                    <span class="n">val_add_node_names</span><span class="p">,</span>
                    <span class="n">model_file_location</span><span class="p">,</span>
                    <span class="n">node_attr_to_name</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="c1"># insert function_proto into model_proto.</span>
            <span class="n">proto</span> <span class="o">=</span> <span class="n">onnx_proto_utils</span><span class="o">.</span><span class="n">_add_onnxscript_fn</span><span class="p">(</span>
                <span class="n">proto</span><span class="p">,</span>
                <span class="n">custom_opsets</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
                <span class="n">_C</span><span class="o">.</span><span class="n">_jit_onnx_log</span><span class="p">(</span><span class="s2">&quot;Exported graph: &quot;</span><span class="p">,</span> <span class="n">graph</span><span class="p">)</span>
            <span class="n">onnx_proto_utils</span><span class="o">.</span><span class="n">_export_file</span><span class="p">(</span><span class="n">proto</span><span class="p">,</span> <span class="n">f</span><span class="p">,</span> <span class="n">export_map</span><span class="p">)</span>
    <span class="k">finally</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">in_onnx_export</span>
        <span class="n">GLOBALS</span><span class="o">.</span><span class="n">in_onnx_export</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="n">GLOBALS</span><span class="o">.</span><span class="n">autograd_inlining</span> <span class="o">=</span> <span class="n">_autograd_inlining_previous</span>
        <span class="n">_reset_trace_module_map</span><span class="p">()</span>

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


<span class="k">def</span> <span class="nf">_apply_friendly_debug_names</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">params</span><span class="p">):</span>
    <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">graph</span><span class="o">.</span><span class="n">nodes</span><span class="p">():</span>
        <span class="k">for</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">n</span><span class="o">.</span><span class="n">inputs</span><span class="p">():</span>
            <span class="n">old_name</span> <span class="o">=</span> <span class="n">v</span><span class="o">.</span><span class="n">debugName</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">old_name</span> <span class="o">!=</span> <span class="nb">str</span><span class="p">(</span><span class="n">v</span><span class="o">.</span><span class="n">unique</span><span class="p">()):</span>
                <span class="k">continue</span>
            <span class="n">new_name</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">n</span><span class="o">.</span><span class="n">kind</span><span class="p">()</span><span class="si">}</span><span class="s2">_</span><span class="si">{</span><span class="n">v</span><span class="o">.</span><span class="n">unique</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span>
            <span class="n">v</span><span class="o">.</span><span class="n">setDebugName</span><span class="p">(</span><span class="n">new_name</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">old_name</span> <span class="ow">in</span> <span class="n">params</span><span class="p">:</span>
                <span class="n">params</span><span class="p">[</span><span class="n">new_name</span><span class="p">]</span> <span class="o">=</span> <span class="n">params</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="n">old_name</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_set_input_and_output_names</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">input_names</span><span class="p">,</span> <span class="n">output_names</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">set_names</span><span class="p">(</span><span class="n">node_list</span><span class="p">,</span> <span class="n">name_list</span><span class="p">,</span> <span class="n">descriptor</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">name_list</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">return</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">name_list</span><span class="p">)</span> <span class="o">&gt;</span> <span class="nb">len</span><span class="p">(</span><span class="n">node_list</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
                <span class="s2">&quot;number of </span><span class="si">%s</span><span class="s2"> names provided (</span><span class="si">%d</span><span class="s2">) exceeded number of </span><span class="si">%s</span><span class="s2">s (</span><span class="si">%d</span><span class="s2">)&quot;</span>
                <span class="o">%</span> <span class="p">(</span><span class="n">descriptor</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">name_list</span><span class="p">),</span> <span class="n">descriptor</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">node_list</span><span class="p">))</span>
            <span class="p">)</span>

        <span class="c1"># Mark if the output node DebugName is set before.</span>
        <span class="n">output_node_set</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">node</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">zip</span><span class="p">(</span><span class="n">name_list</span><span class="p">,</span> <span class="n">node_list</span><span class="p">)):</span>
            <span class="c1"># Duplicated output node, insert onnx::Identity to avoid setting the same DebugName after setDebugName().</span>
            <span class="k">if</span> <span class="n">descriptor</span> <span class="o">==</span> <span class="s2">&quot;output&quot;</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">output_node_set</span><span class="p">:</span>
                    <span class="n">identity_node</span> <span class="o">=</span> <span class="n">graph</span><span class="o">.</span><span class="n">create</span><span class="p">(</span><span class="s2">&quot;onnx::Identity&quot;</span><span class="p">)</span>
                    <span class="n">identity_node</span><span class="o">.</span><span class="n">insertAfter</span><span class="p">(</span><span class="n">node</span><span class="o">.</span><span class="n">node</span><span class="p">())</span>
                    <span class="n">identity_node</span><span class="o">.</span><span class="n">addInput</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>
                    <span class="n">identity_node</span><span class="o">.</span><span class="n">output</span><span class="p">()</span><span class="o">.</span><span class="n">setType</span><span class="p">(</span><span class="n">node</span><span class="o">.</span><span class="n">type</span><span class="p">())</span>
                    <span class="n">graph</span><span class="o">.</span><span class="n">return_node</span><span class="p">()</span><span class="o">.</span><span class="n">replaceInput</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">identity_node</span><span class="o">.</span><span class="n">output</span><span class="p">())</span>
                    <span class="n">node</span> <span class="o">=</span> <span class="n">identity_node</span><span class="o">.</span><span class="n">output</span><span class="p">()</span>
                <span class="n">output_node_set</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">node</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">node</span><span class="o">.</span><span class="n">debugName</span><span class="p">()</span> <span class="o">!=</span> <span class="n">name</span><span class="p">:</span>
                <span class="n">node</span><span class="o">.</span><span class="n">setDebugName</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>

    <span class="n">set_names</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">graph</span><span class="o">.</span><span class="n">inputs</span><span class="p">()),</span> <span class="n">input_names</span><span class="p">,</span> <span class="s2">&quot;input&quot;</span><span class="p">)</span>
    <span class="n">set_names</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">graph</span><span class="o">.</span><span class="n">outputs</span><span class="p">()),</span> <span class="n">output_names</span><span class="p">,</span> <span class="s2">&quot;output&quot;</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_run_symbolic_method</span><span class="p">(</span><span class="n">g</span><span class="p">,</span> <span class="n">op_name</span><span class="p">,</span> <span class="n">symbolic_fn</span><span class="p">,</span> <span class="n">args</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This trampoline function gets invoked for every symbolic method</span>
<span class="sd">    call from C++.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">graph_context</span> <span class="o">=</span> <span class="n">jit_utils</span><span class="o">.</span><span class="n">GraphContext</span><span class="p">(</span>
            <span class="n">graph</span><span class="o">=</span><span class="n">g</span><span class="p">,</span>
            <span class="n">block</span><span class="o">=</span><span class="n">g</span><span class="o">.</span><span class="n">block</span><span class="p">(),</span>
            <span class="n">opset</span><span class="o">=</span><span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span><span class="p">,</span>
            <span class="n">original_node</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>  <span class="c1"># type: ignore[arg-type]</span>
            <span class="n">params_dict</span><span class="o">=</span><span class="n">_params_dict</span><span class="p">,</span>
            <span class="n">env</span><span class="o">=</span><span class="p">{},</span>
            <span class="n">values_in_env</span><span class="o">=</span><span class="nb">set</span><span class="p">(),</span>
            <span class="n">new_nodes</span><span class="o">=</span><span class="p">[],</span>
        <span class="p">)</span>
        <span class="k">return</span> <span class="n">symbolic_fn</span><span class="p">(</span><span class="n">graph_context</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">)</span>
    <span class="k">except</span> <span class="ne">TypeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="c1"># Handle the specific case where we didn&#39;t successfully dispatch</span>
        <span class="c1"># to symbolic_fn.  Otherwise, the backtrace will have the clues</span>
        <span class="c1"># you need.</span>
        <span class="n">e</span><span class="o">.</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">e</span><span class="o">.</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2"> (occurred when translating </span><span class="si">{</span><span class="n">op_name</span><span class="si">}</span><span class="s2">)&quot;</span><span class="p">,)</span>
        <span class="k">raise</span>


<span class="k">def</span> <span class="nf">_add_block</span><span class="p">(</span><span class="n">node</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Node</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">_C</span><span class="o">.</span><span class="n">Block</span><span class="p">:</span>
    <span class="k">return</span> <span class="n">node</span><span class="o">.</span><span class="n">addBlock</span><span class="p">()</span>


<span class="k">def</span> <span class="nf">_add_input_to_block</span><span class="p">(</span><span class="n">block</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Block</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">block</span><span class="o">.</span><span class="n">addInputToBlock</span><span class="p">()</span>  <span class="c1"># type: ignore[attr-defined]</span>


<span class="k">def</span> <span class="nf">_add_output_to_block</span><span class="p">(</span><span class="n">block</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Block</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Value</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</span><span class="p">:</span>
    <span class="k">return</span> <span class="n">block</span><span class="o">.</span><span class="n">registerOutput</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_should_aten_fallback</span><span class="p">(</span>
    <span class="n">name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">opset_version</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">:</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span>
<span class="p">):</span>
    <span class="c1"># For all builds, if domain==&quot;aten&quot; and operator_export_type==ONNX_ATEN,</span>
    <span class="c1">#   an aten::ATen operator is created regardless of symbolics existence</span>

    <span class="n">is_exportable_aten_op</span> <span class="o">=</span> <span class="n">registration</span><span class="o">.</span><span class="n">registry</span><span class="o">.</span><span class="n">is_registered_op</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">opset_version</span><span class="p">)</span>
    <span class="n">is_onnx_aten_export</span> <span class="o">=</span> <span class="n">operator_export_type</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX_ATEN</span>
    <span class="n">is_aten_fallback_export</span> <span class="o">=</span> <span class="p">(</span>
        <span class="n">operator_export_type</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX_ATEN_FALLBACK</span>
    <span class="p">)</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="n">name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;aten::&quot;</span><span class="p">):</span>
        <span class="k">return</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">is_onnx_aten_export</span> <span class="ow">or</span> <span class="p">(</span><span class="n">is_aten_fallback_export</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">is_exportable_aten_op</span><span class="p">):</span>
        <span class="k">return</span> <span class="kc">True</span>

    <span class="k">return</span> <span class="kc">False</span>


<span class="k">def</span> <span class="nf">_get_aten_op_overload_name</span><span class="p">(</span><span class="n">n</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Node</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
    <span class="c1"># Returns `overload_name` attribute to ATen ops on non-Caffe2 builds</span>
    <span class="n">schema</span> <span class="o">=</span> <span class="n">n</span><span class="o">.</span><span class="n">schema</span><span class="p">()</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">schema</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;aten::&quot;</span><span class="p">):</span>
        <span class="k">return</span> <span class="s2">&quot;&quot;</span>
    <span class="k">return</span> <span class="n">_C</span><span class="o">.</span><span class="n">parse_schema</span><span class="p">(</span><span class="n">schema</span><span class="p">)</span><span class="o">.</span><span class="n">overload_name</span>


<span class="k">def</span> <span class="nf">_run_symbolic_function</span><span class="p">(</span>
    <span class="n">graph</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Graph</span><span class="p">,</span>
    <span class="n">block</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Block</span><span class="p">,</span>
    <span class="n">node</span><span class="p">:</span> <span class="n">_C</span><span class="o">.</span><span class="n">Node</span><span class="p">,</span>
    <span class="n">inputs</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span>
    <span class="n">env</span><span class="p">:</span> <span class="nb">dict</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">Value</span><span class="p">,</span> <span class="n">_C</span><span class="o">.</span><span class="n">Value</span><span class="p">],</span>
    <span class="n">values_in_env</span><span class="p">:</span> <span class="nb">set</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">Value</span><span class="p">],</span>
    <span class="n">new_nodes</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">Node</span><span class="p">],</span>
    <span class="n">operator_export_type</span><span class="o">=</span><span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">_C</span><span class="o">.</span><span class="n">Value</span> <span class="o">|</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">_C</span><span class="o">.</span><span class="n">Value</span> <span class="o">|</span> <span class="kc">None</span><span class="p">]</span> <span class="o">|</span> <span class="kc">None</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Runs a symbolic function.</span>

<span class="sd">    The function is used in C++ to export the node to ONNX.</span>

<span class="sd">    Returns:</span>
<span class="sd">        A single or a tuple of Values.</span>
<span class="sd">        None when the node gets cloned as is into the new graph.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">opset_version</span> <span class="o">=</span> <span class="n">GLOBALS</span><span class="o">.</span><span class="n">export_onnx_opset_version</span>

    <span class="c1"># See Note [Export inplace]</span>
    <span class="n">node_kind</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">kind</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">node_kind</span><span class="o">.</span><span class="n">endswith</span><span class="p">(</span><span class="s2">&quot;_&quot;</span><span class="p">):</span>
        <span class="c1"># Treat relu_ -&gt; relu; add_ -&gt; add etc.</span>
        <span class="n">ns_op_name</span> <span class="o">=</span> <span class="n">node_kind</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ns_op_name</span> <span class="o">=</span> <span class="n">node_kind</span>

    <span class="n">namespace</span><span class="p">,</span> <span class="n">op_name</span> <span class="o">=</span> <span class="n">jit_utils</span><span class="o">.</span><span class="n">parse_node_kind</span><span class="p">(</span><span class="n">ns_op_name</span><span class="p">)</span>

    <span class="n">graph_context</span> <span class="o">=</span> <span class="n">jit_utils</span><span class="o">.</span><span class="n">GraphContext</span><span class="p">(</span>
        <span class="n">graph</span><span class="o">=</span><span class="n">graph</span><span class="p">,</span>
        <span class="n">block</span><span class="o">=</span><span class="n">block</span><span class="p">,</span>
        <span class="n">opset</span><span class="o">=</span><span class="n">opset_version</span><span class="p">,</span>
        <span class="n">original_node</span><span class="o">=</span><span class="n">node</span><span class="p">,</span>
        <span class="n">params_dict</span><span class="o">=</span><span class="n">_params_dict</span><span class="p">,</span>
        <span class="n">env</span><span class="o">=</span><span class="n">env</span><span class="p">,</span>
        <span class="n">values_in_env</span><span class="o">=</span><span class="n">values_in_env</span><span class="p">,</span>
        <span class="n">new_nodes</span><span class="o">=</span><span class="n">new_nodes</span><span class="p">,</span>
    <span class="p">)</span>

    <span class="c1"># Direct ATen export requested</span>
    <span class="k">if</span> <span class="n">_should_aten_fallback</span><span class="p">(</span><span class="n">ns_op_name</span><span class="p">,</span> <span class="n">opset_version</span><span class="p">,</span> <span class="n">operator_export_type</span><span class="p">):</span>
        <span class="n">attrs</span> <span class="o">=</span> <span class="p">{</span>
            <span class="n">k</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="n">node</span><span class="o">.</span><span class="n">kindOf</span><span class="p">(</span><span class="n">k</span><span class="p">)[</span><span class="mi">0</span><span class="p">]:</span> <span class="n">symbolic_helper</span><span class="o">.</span><span class="n">_node_get</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">k</span><span class="p">)</span>
            <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">attributeNames</span><span class="p">()</span>
        <span class="p">}</span>
        <span class="n">outputs</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">outputsSize</span><span class="p">()</span>
        <span class="n">attrs</span><span class="p">[</span><span class="s2">&quot;outputs&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">outputs</span>
        <span class="k">return</span> <span class="n">graph_context</span><span class="o">.</span><span class="n">aten_op</span><span class="p">(</span>
            <span class="n">op_name</span><span class="p">,</span>
            <span class="o">*</span><span class="n">inputs</span><span class="p">,</span>
            <span class="n">overload_name</span><span class="o">=</span><span class="n">_get_aten_op_overload_name</span><span class="p">(</span><span class="n">node</span><span class="p">),</span>
            <span class="o">**</span><span class="n">attrs</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="k">try</span><span class="p">:</span>
        <span class="n">domain</span> <span class="o">=</span> <span class="n">namespace</span>
        <span class="n">symbolic_function_name</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">domain</span><span class="si">}</span><span class="s2">::</span><span class="si">{</span><span class="n">op_name</span><span class="si">}</span><span class="s2">&quot;</span>

        <span class="n">symbolic_function_group</span> <span class="o">=</span> <span class="n">registration</span><span class="o">.</span><span class="n">registry</span><span class="o">.</span><span class="n">get_function_group</span><span class="p">(</span>
            <span class="n">symbolic_function_name</span>
        <span class="p">)</span>
        <span class="k">if</span> <span class="n">symbolic_function_group</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">symbolic_fn</span> <span class="o">=</span> <span class="n">symbolic_function_group</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">opset_version</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">symbolic_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># TODO Wrap almost identical attrs assignment or comment the difference.</span>
                <span class="n">attrs</span> <span class="o">=</span> <span class="p">{</span>
                    <span class="n">k</span><span class="p">:</span> <span class="n">symbolic_helper</span><span class="o">.</span><span class="n">_node_get</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">k</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">attributeNames</span><span class="p">()</span>
                <span class="p">}</span>
                <span class="k">return</span> <span class="n">symbolic_fn</span><span class="p">(</span><span class="n">graph_context</span><span class="p">,</span> <span class="o">*</span><span class="n">inputs</span><span class="p">,</span> <span class="o">**</span><span class="n">attrs</span><span class="p">)</span>

        <span class="n">attrs</span> <span class="o">=</span> <span class="p">{</span>
            <span class="n">k</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="n">node</span><span class="o">.</span><span class="n">kindOf</span><span class="p">(</span><span class="n">k</span><span class="p">)[</span><span class="mi">0</span><span class="p">]:</span> <span class="n">symbolic_helper</span><span class="o">.</span><span class="n">_node_get</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">k</span><span class="p">)</span>
            <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">attributeNames</span><span class="p">()</span>
        <span class="p">}</span>
        <span class="k">if</span> <span class="n">namespace</span> <span class="o">==</span> <span class="s2">&quot;onnx&quot;</span><span class="p">:</span>
            <span class="c1"># Clone node to trigger ONNX shape inference</span>
            <span class="k">return</span> <span class="n">graph_context</span><span class="o">.</span><span class="n">op</span><span class="p">(</span>
                <span class="n">op_name</span><span class="p">,</span> <span class="o">*</span><span class="n">inputs</span><span class="p">,</span> <span class="o">**</span><span class="n">attrs</span><span class="p">,</span> <span class="n">outputs</span><span class="o">=</span><span class="n">node</span><span class="o">.</span><span class="n">outputsSize</span><span class="p">()</span>
            <span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>

        <span class="k">raise</span> <span class="n">errors</span><span class="o">.</span><span class="n">UnsupportedOperatorError</span><span class="p">(</span>
            <span class="n">symbolic_function_name</span><span class="p">,</span>
            <span class="n">opset_version</span><span class="p">,</span>
            <span class="n">symbolic_function_group</span><span class="o">.</span><span class="n">get_min_supported</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">symbolic_function_group</span>
            <span class="k">else</span> <span class="kc">None</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="k">except</span> <span class="ne">RuntimeError</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">operator_export_type</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX_FALLTHROUGH</span><span class="p">:</span>
            <span class="k">return</span> <span class="kc">None</span>
        <span class="k">elif</span> <span class="n">operator_export_type</span> <span class="o">==</span> <span class="n">_C_onnx</span><span class="o">.</span><span class="n">OperatorExportTypes</span><span class="o">.</span><span class="n">ONNX_ATEN_FALLBACK</span><span class="p">:</span>
            <span class="c1"># Emit ATen op for non-Caffe2 builds when `operator_export_type==ONNX_ATEN_FALLBACK`</span>
            <span class="n">attrs</span> <span class="o">=</span> <span class="p">{</span>
                <span class="n">k</span> <span class="o">+</span> <span class="s2">&quot;_&quot;</span> <span class="o">+</span> <span class="n">node</span><span class="o">.</span><span class="n">kindOf</span><span class="p">(</span><span class="n">k</span><span class="p">)[</span><span class="mi">0</span><span class="p">]:</span> <span class="n">symbolic_helper</span><span class="o">.</span><span class="n">_node_get</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">k</span><span class="p">)</span>
                <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">node</span><span class="o">.</span><span class="n">attributeNames</span><span class="p">()</span>
            <span class="p">}</span>
            <span class="k">return</span> <span class="n">graph_context</span><span class="o">.</span><span class="n">aten_op</span><span class="p">(</span>
                <span class="n">op_name</span><span class="p">,</span>
                <span class="o">*</span><span class="n">inputs</span><span class="p">,</span>
                <span class="n">overload_name</span><span class="o">=</span><span class="n">_get_aten_op_overload_name</span><span class="p">(</span><span class="n">node</span><span class="p">),</span>
                <span class="o">**</span><span class="n">attrs</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="k">raise</span>
    <span class="k">except</span> <span class="ne">TypeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="c1"># Handle the specific case where we didn&#39;t successfully dispatch.</span>
        <span class="c1"># Otherwise, the backtrace will have the clues you need.</span>
        <span class="n">e</span><span class="o">.</span><span class="n">args</span> <span class="o">=</span> <span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">e</span><span class="o">.</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2"> </span><span class="se">\n</span><span class="s2">(Occurred when translating </span><span class="si">{</span><span class="n">op_name</span><span class="si">}</span><span class="s2">).&quot;</span><span class="p">,)</span>
        <span class="k">raise</span>


<span class="k">def</span> <span class="nf">_verify_custom_op_name</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">re</span><span class="o">.</span><span class="n">match</span><span class="p">(</span><span class="sa">r</span><span class="s2">&quot;^[a-zA-Z0-9-_]+::[a-zA-Z-_]+[a-zA-Z0-9-_]*$&quot;</span><span class="p">,</span> <span class="n">symbolic_name</span><span class="p">):</span>
        <span class="k">raise</span> <span class="n">errors</span><span class="o">.</span><span class="n">OnnxExporterError</span><span class="p">(</span>
            <span class="sa">f</span><span class="s2">&quot;Failed to register operator </span><span class="si">{</span><span class="n">symbolic_name</span><span class="si">}</span><span class="s2">. &quot;</span>
            <span class="s2">&quot;The symbolic name must match the format domain::name, &quot;</span>
            <span class="s2">&quot;and should start with a letter and contain only &quot;</span>
            <span class="s2">&quot;alphanumerical characters&quot;</span>
        <span class="p">)</span>

    <span class="n">ns</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">jit_utils</span><span class="o">.</span><span class="n">parse_node_kind</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">ns</span> <span class="o">==</span> <span class="s2">&quot;onnx&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="sa">f</span><span class="s2">&quot;Failed to register operator </span><span class="si">{</span><span class="n">symbolic_name</span><span class="si">}</span><span class="s2">. </span><span class="si">{</span><span class="n">ns</span><span class="si">}</span><span class="s2"> domain cannot be modified.&quot;</span>
        <span class="p">)</span>


<div class="viewcode-block" id="register_custom_op_symbolic"><a class="viewcode-back" href="../../../onnx_torchscript.html#torch.onnx.register_custom_op_symbolic">[docs]</a><span class="k">def</span> <span class="nf">register_custom_op_symbolic</span><span class="p">(</span>
    <span class="n">symbolic_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
    <span class="n">symbolic_fn</span><span class="p">:</span> <span class="n">Callable</span><span class="p">,</span>
    <span class="n">opset_version</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Registers a symbolic function for a custom operator.</span>

<span class="sd">    When the user registers symbolic for custom/contrib ops,</span>
<span class="sd">    it is highly recommended to add shape inference for that operator via setType API,</span>
<span class="sd">    otherwise the exported graph may have incorrect shape inference in some extreme cases.</span>
<span class="sd">    An example of setType is `test_aten_embedding_2` in `test_operators.py`.</span>

<span class="sd">    See &quot;Custom Operators&quot; in the module documentation for an example usage.</span>

<span class="sd">    Args:</span>
<span class="sd">        symbolic_name (str): The name of the custom operator in &quot;&lt;domain&gt;::&lt;op&gt;&quot;</span>
<span class="sd">            format.</span>
<span class="sd">        symbolic_fn (Callable): A function that takes in the ONNX graph and</span>
<span class="sd">            the input arguments to the current operator, and returns new</span>
<span class="sd">            operator nodes to add to the graph.</span>
<span class="sd">        opset_version (int): The ONNX opset version in which to register.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">symbolic_name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;::&quot;</span><span class="p">):</span>
        <span class="n">symbolic_name</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;aten</span><span class="si">{</span><span class="n">symbolic_name</span><span class="si">}</span><span class="s2">&quot;</span>

    <span class="n">_verify_custom_op_name</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">)</span>

    <span class="n">registration</span><span class="o">.</span><span class="n">custom_onnx_symbolic</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">,</span> <span class="n">opset_version</span><span class="p">)(</span><span class="n">symbolic_fn</span><span class="p">)</span></div>


<div class="viewcode-block" id="unregister_custom_op_symbolic"><a class="viewcode-back" href="../../../onnx_torchscript.html#torch.onnx.unregister_custom_op_symbolic">[docs]</a><span class="k">def</span> <span class="nf">unregister_custom_op_symbolic</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">opset_version</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Unregisters ``symbolic_name``.</span>

<span class="sd">    See &quot;Custom Operators&quot; in the module documentation for an example usage.</span>

<span class="sd">    Args:</span>
<span class="sd">        symbolic_name (str): The name of the custom operator in &quot;&lt;domain&gt;::&lt;op&gt;&quot;</span>
<span class="sd">            format.</span>
<span class="sd">        opset_version (int): The ONNX opset version in which to unregister.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">symbolic_name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;::&quot;</span><span class="p">):</span>
        <span class="n">symbolic_name</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;aten</span><span class="si">{</span><span class="n">symbolic_name</span><span class="si">}</span><span class="s2">&quot;</span>

    <span class="n">_verify_custom_op_name</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">)</span>

    <span class="n">registration</span><span class="o">.</span><span class="n">registry</span><span class="o">.</span><span class="n">unregister</span><span class="p">(</span><span class="n">symbolic_name</span><span class="p">,</span> <span class="n">opset_version</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_validate_dynamic_axes</span><span class="p">(</span><span class="n">dynamic_axes</span><span class="p">,</span> <span class="n">model</span><span class="p">,</span> <span class="n">input_names</span><span class="p">,</span> <span class="n">output_names</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Ensures dynamic axes argument is follows the expected format.&quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">dynamic_axes</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;graph&quot;</span><span class="p">):</span>
        <span class="c1"># Extracting set of valid input/output names that shall be used for dynamic_axes</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">input_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">input_names</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">input_names</span> <span class="o">=</span> <span class="p">[</span><span class="n">x</span><span class="o">.</span><span class="n">debugName</span><span class="p">()</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">graph</span><span class="o">.</span><span class="n">inputs</span><span class="p">()]</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">output_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">output_names</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">output_names</span> <span class="o">=</span> <span class="p">[</span><span class="n">y</span><span class="o">.</span><span class="n">debugName</span><span class="p">()</span> <span class="k">for</span> <span class="n">y</span> <span class="ow">in</span> <span class="n">model</span><span class="o">.</span><span class="n">graph</span><span class="o">.</span><span class="n">outputs</span><span class="p">()]</span>

    <span class="n">valid_names</span> <span class="o">=</span> <span class="nb">set</span><span class="p">((</span><span class="n">input_names</span> <span class="ow">or</span> <span class="p">[])</span> <span class="o">+</span> <span class="p">(</span><span class="n">output_names</span> <span class="ow">or</span> <span class="p">[]))</span>

    <span class="c1"># If dynamic axes are provided as a list rather than dictionary, they should</span>
    <span class="c1"># first get converted to a dictionary in expected format. If desired axes names</span>
    <span class="c1"># are not provided for dynamic axes, automatic names shall be generated for</span>
    <span class="c1"># provided dynamic axes of specified input/output</span>
    <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">dynamic_axes</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">valid_names</span><span class="p">:</span>
            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
                <span class="sa">f</span><span class="s2">&quot;Provided key </span><span class="si">{</span><span class="n">key</span><span class="si">}</span><span class="s2"> for dynamic axes is not a valid input/output name&quot;</span>
            <span class="p">)</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
                <span class="s2">&quot;No names were found for specified dynamic axes of provided input.&quot;</span>
                <span class="sa">f</span><span class="s2">&quot;Automatically generated names will be applied to each dynamic axes of input </span><span class="si">{</span><span class="n">key</span><span class="si">}</span><span class="s2">&quot;</span>
            <span class="p">)</span>

            <span class="n">value_dict</span> <span class="o">=</span> <span class="p">{}</span>
            <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">value</span><span class="p">):</span>
                <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                        <span class="s2">&quot;The type of axis index is expected to be an integer&quot;</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">value_dict</span><span class="p">:</span>
                    <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span>
                        <span class="sa">f</span><span class="s2">&quot;Duplicate dynamic axis index </span><span class="si">{</span><span class="n">x</span><span class="si">}</span><span class="s2"> was provided for input </span><span class="si">{</span><span class="n">key</span><span class="si">}</span><span class="s2">.&quot;</span>
                    <span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">value_dict</span><span class="p">[</span><span class="n">x</span><span class="p">]</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">key</span><span class="p">)</span> <span class="o">+</span> <span class="s2">&quot;_dynamic_axes_&quot;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">dynamic_axes</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">value_dict</span>


<span class="k">def</span> <span class="nf">model_signature</span><span class="p">(</span><span class="n">model</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="o">|</span> <span class="n">Callable</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">inspect</span><span class="o">.</span><span class="n">Signature</span><span class="p">:</span>
    <span class="k">return</span> <span class="n">inspect</span><span class="o">.</span><span class="n">signature</span><span class="p">(</span>
        <span class="n">model</span><span class="o">.</span><span class="n">forward</span> <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model</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">else</span> <span class="n">model</span>
    <span class="p">)</span>
</pre></div>

             </article>
             
            </div>
            <footer>
  

    <hr>

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

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

</footer>

          </div>
<script>

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

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

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

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

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

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

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


  <!-- Begin Footer -->

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

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

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

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

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

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

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

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

  </footer>

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

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

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

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

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

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

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

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

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

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

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

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

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

  <!-- End Mobile Menu -->

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

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

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