2.0/dynamo/custom-backends.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="generator" content="Docutils 0.18.1: http://docutils.sourceforge.net/" />

  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>Custom Backends &mdash; PyTorch 2.0 documentation</title>
  

    <link rel="canonical" href="https://pytorch.org/docs/stable/torch.compiler_custom_backends.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="index" title="Index" href="../genindex.html" />
    <link rel="search" title="Search" href="../search.html" />
    <link rel="next" title="TorchDynamo Deeper Dive" href="deep-dive.html" />
    <link rel="prev" title="Guards Overview" href="guards-overview.html" />


  <!-- Google Analytics -->
  
    <script async src="https://www.googletagmanager.com/gtag/js?id=UA-117752657-2"></script>
    <script>
      window.dataLayer = window.dataLayer || [];
      function gtag(){dataLayer.push(arguments);}
      gtag('js', new Date());

      gtag('config', 'UA-117752657-2');
    </script>
  
  <!-- End Google Analytics -->
  

  <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>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active docs-active">
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="resource-option with-down-orange-arrow">
                Docs
              </a>
              <div class="resources-dropdown-menu">
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/docs/stable/index.html">
                  <span class="dropdown-title">PyTorch</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/audio/stable/index.html">
                  <span class="dropdown-title">torchaudio</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/text/stable/index.html">
                  <span class="dropdown-title">torchtext</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/vision/stable/index.html">
                  <span class="dropdown-title">torchvision</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torcharrow">
                  <span class="dropdown-title">torcharrow</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/data">
                  <span class="dropdown-title">TorchData</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torchrec">
                  <span class="dropdown-title">TorchRec</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/serve/">
                  <span class="dropdown-title">TorchServe</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torchx/">
                  <span class="dropdown-title">TorchX</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/xla">
                  <span class="dropdown-title">PyTorch on XLA Devices</span>
                  <p></p>
                </a>
            </div>
          </li>

          <li>
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="resource-option with-down-arrow">
                Resources
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/features">
                  <span class="dropdown-title">About</span>
                  <p>Learn about PyTorch’s features and capabilities</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/foundation">
                  <span class="dropdown-title">PyTorch Foundation</span>
                  <p>Learn about the PyTorch foundation</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://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/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/events">
                  <span class="dropdown-title">Events</span>
                  <p>Find events, webinars, and podcasts</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/hub">
                  <span class="dropdown-title">Models (Beta)</span>
                  <p>Discover, publish, and reuse pre-trained models</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">GitHub</a>
          </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.0 &#x25BC</a>
    </div>
    

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

          
              <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">CUDA 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/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/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">torch.compile</span></p>
<ul class="current">
<li class="toctree-l1 current"><a class="reference internal" href="index.html">TorchDynamo Overview</a></li>
<li class="toctree-l1"><a class="reference internal" href="installation.html">Installing TorchDynamo</a></li>
<li class="toctree-l1"><a class="reference internal" href="get-started.html">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="guards-overview.html">Guards Overview</a></li>
<li class="toctree-l1 current"><a class="current reference internal" href="#">Custom Backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="deep-dive.html">TorchDynamo Deeper Dive</a></li>
<li class="toctree-l1"><a class="reference internal" href="troubleshooting.html">TorchDynamo Troubleshooting</a></li>
<li class="toctree-l1"><a class="reference internal" href="faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../ir.html">IRs</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="../cuda.html">torch.cuda</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="../backends.html">torch.backends</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.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.optim.html">torch.distributed.optim</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.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="../_dynamo.html">torch._dynamo</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="../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="../onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../onnx_diagnostics.html">torch.onnx diagnostics</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="../pipeline.html">Pipeline Parallelism</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="../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="../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="../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="../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>
</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">TorchDynamo Overview</a> &gt;</li>
        
      <li>Custom Backends</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
            
            <a href="../_sources/dynamo/custom-backends.rst.txt" rel="nofollow"><img src="../_static/images/view-page-source-icon.svg"></a>
          
        
      </li>
    
  </ul>

  
</div>
        </div>

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

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

        
          <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">
              
  <section id="custom-backends">
<h1>Custom Backends<a class="headerlink" href="#custom-backends" title="Permalink to this heading">¶</a></h1>
<section id="overview">
<h2>Overview<a class="headerlink" href="#overview" title="Permalink to this heading">¶</a></h2>
<p><code class="docutils literal notranslate"><span class="pre">torch.compile</span></code> provides a straightforward method to enable users
to define custom backends.</p>
<p>A backend function has the contract
<code class="docutils literal notranslate"><span class="pre">(gm:</span> <span class="pre">torch.fx.GraphModule,</span> <span class="pre">example_inputs:</span> <span class="pre">List[torch.Tensor])</span> <span class="pre">-&gt;</span> <span class="pre">Callable</span></code>.</p>
<p>Backend functions can be called by TorchDynamo, the graph tracing component of <code class="docutils literal notranslate"><span class="pre">torch.compile</span></code>,
after tracing an FX graph and are
expected to return a compiled function that is equivalent to the traced FX graph.
The returned callable should have the same contract as the <code class="docutils literal notranslate"><span class="pre">forward</span></code> function of the original <code class="docutils literal notranslate"><span class="pre">torch.fx.GraphModule</span></code>
passed into the backend:
<code class="docutils literal notranslate"><span class="pre">(*args:</span> <span class="pre">torch.Tensor)</span> <span class="pre">-&gt;</span> <span class="pre">List[torch.Tensor]</span></code>.</p>
<p>In order for TorchDynamo to call your backend, pass your backend function as the <code class="docutils literal notranslate"><span class="pre">backend</span></code> kwarg in
<code class="docutils literal notranslate"><span class="pre">torch.compile</span></code>. For example,</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">torch</span>

<span class="k">def</span> <span class="nf">my_custom_backend</span><span class="p">(</span><span class="n">gm</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">gm</span><span class="o">.</span><span class="n">forward</span>

<span class="k">def</span> <span class="nf">f</span><span class="p">(</span><span class="o">...</span><span class="p">):</span>
    <span class="o">...</span>

<span class="n">f_opt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="n">backend</span><span class="o">=</span><span class="n">my_custom_backend</span><span class="p">)</span>

<span class="nd">@torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">backend</span><span class="o">=</span><span class="n">my_custom_backend</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">g</span><span class="p">(</span><span class="o">...</span><span class="p">):</span>
    <span class="o">...</span>
</pre></div>
</div>
<p>See below for more examples.</p>
</section>
<section id="registering-custom-backends">
<h2>Registering Custom Backends<a class="headerlink" href="#registering-custom-backends" title="Permalink to this heading">¶</a></h2>
<p>You can register your backend using the <code class="docutils literal notranslate"><span class="pre">register_backend</span></code> decorator, for example,</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">torch._dynamo.optimizations</span> <span class="kn">import</span> <span class="n">register_backend</span>

<span class="nd">@register_backend</span>
<span class="k">def</span> <span class="nf">my_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">):</span>
    <span class="o">...</span>
</pre></div>
</div>
<p>Besides the <code class="docutils literal notranslate"><span class="pre">register_backend</span></code> decorator, if your backend is in another python package, you could also register your
backend through entry points of python package, which provides a way for a package to register a plugin for another one.</p>
<div class="admonition hint">
<p class="admonition-title">Hint</p>
<p>You can learn more about <code class="docutils literal notranslate"><span class="pre">entry_points</span></code> in the
<a class="reference external" href="https://setuptools.pypa.io/en/latest/userguide/entry_point.html">python packaging documentation</a>.</p>
</div>
<p>To register your backend through <code class="docutils literal notranslate"><span class="pre">entry_points</span></code>, you could add your backend function to the <code class="docutils literal notranslate"><span class="pre">torch_dynamo_backends</span></code> entry point group in the
<code class="docutils literal notranslate"><span class="pre">setup.py</span></code> file of your package like:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="o">...</span>
<span class="n">setup</span><span class="p">(</span>
    <span class="o">...</span>
    <span class="s1">&#39;torch_dynamo_backends&#39;</span><span class="p">:</span> <span class="p">[</span>
        <span class="s1">&#39;my_compiler = your_module.submodule:my_compiler&#39;</span><span class="p">,</span>
    <span class="p">]</span>
    <span class="o">...</span>
<span class="p">)</span>
</pre></div>
</div>
<p>Please replace the <code class="docutils literal notranslate"><span class="pre">my_compiler</span></code> before <code class="docutils literal notranslate"><span class="pre">=</span></code> to the name of your backend’s name and replace the part after <code class="docutils literal notranslate"><span class="pre">=</span></code> to
the module and function name of your backend function.
The entry point will be added to your python environment after the installation of the package.
When you call <code class="docutils literal notranslate"><span class="pre">torch.compile(model,</span> <span class="pre">backend=&quot;my_compiler&quot;)</span></code>, PyTorch would first search the backend named <code class="docutils literal notranslate"><span class="pre">my_compiler</span></code>
that has been registered with <code class="docutils literal notranslate"><span class="pre">register_backend</span></code>. If not found, it will continue to search in all backends registered
via <code class="docutils literal notranslate"><span class="pre">entry_points</span></code>.</p>
<p>Registration serves two purposes:</p>
<ul class="simple">
<li><p>You can pass a string containing your backend function’s name to <code class="docutils literal notranslate"><span class="pre">torch.compile</span></code> instead of the function itself,
for example, <code class="docutils literal notranslate"><span class="pre">torch.compile(model,</span> <span class="pre">backend=&quot;my_compiler&quot;)</span></code>.</p></li>
<li><p>It is required for use with the <a class="reference external" href="https://pytorch.org/docs/stable/torch.compiler_troubleshooting.html">minifier</a>. Any generated
code from the minifier must call your code that registers your backend function, typically through an <code class="docutils literal notranslate"><span class="pre">import</span></code> statement.</p></li>
</ul>
</section>
<section id="custom-backends-after-aotautograd">
<h2>Custom Backends after AOTAutograd<a class="headerlink" href="#custom-backends-after-aotautograd" title="Permalink to this heading">¶</a></h2>
<p>It is possible to define custom backends that are called by AOTAutograd rather than TorchDynamo.
This is useful for 2 main reasons:</p>
<ul class="simple">
<li><p>Users can define backends that support model training, as AOTAutograd can generate the backward graph for compilation.</p></li>
<li><p>AOTAutograd produces FX graphs consisting of <a class="reference external" href="https://pytorch.org/docs/master/ir.html#canonical-aten-ir">canonical Aten ops</a>. As a result,
custom backends only need to support the canonical Aten opset, which is a significantly smaller opset than the entire torch/Aten opset.</p></li>
</ul>
<p>Wrap your backend with
<code class="docutils literal notranslate"><span class="pre">torch._dynamo.optimizations.training.aot_autograd</span></code> and use <code class="docutils literal notranslate"><span class="pre">torch.compile</span></code> with the <code class="docutils literal notranslate"><span class="pre">backend</span></code> kwarg as before.
Backend functions wrapped by <code class="docutils literal notranslate"><span class="pre">aot_autograd</span></code> should have the same contract as before.</p>
<p>Backend functions are passed to <code class="docutils literal notranslate"><span class="pre">aot_autograd</span></code> through the <code class="docutils literal notranslate"><span class="pre">fw_compiler</span></code> (forward compiler)
or <code class="docutils literal notranslate"><span class="pre">bw_compiler</span></code> (backward compiler) kwargs. If <code class="docutils literal notranslate"><span class="pre">bw_compiler</span></code> is not specified, the backward compile function
defaults to the forward compile function.</p>
<p>One caveat is that AOTAutograd requires compiled functions returned by backends to be “boxed”. This can be done by wrapping
the compiled function with <code class="docutils literal notranslate"><span class="pre">functorch.compile.make_boxed_func</span></code>.</p>
<p>For example,</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">torch._dynamo.optimizations.training</span> <span class="kn">import</span> <span class="n">aot_autograd</span>
<span class="kn">from</span> <span class="nn">functorch.compile</span> <span class="kn">import</span> <span class="n">make_boxed_func</span>

<span class="k">def</span> <span class="nf">my_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">make_boxed_func</span><span class="p">(</span><span class="n">gm</span><span class="o">.</span><span class="n">forward</span><span class="p">)</span>

<span class="n">my_backend</span> <span class="o">=</span> <span class="n">aot_autograd</span><span class="p">(</span><span class="n">fw_compiler</span><span class="o">=</span><span class="n">my_compiler</span><span class="p">)</span>  <span class="c1"># bw_compiler=my_compiler</span>

<span class="n">model_opt</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">backend</span><span class="o">=</span><span class="n">my_backend</span><span class="p">)</span>
</pre></div>
</div>
</section>
<section id="examples">
<h2>Examples<a class="headerlink" href="#examples" title="Permalink to this heading">¶</a></h2>
<section id="debugging-backend">
<h3>Debugging Backend<a class="headerlink" href="#debugging-backend" title="Permalink to this heading">¶</a></h3>
<p>If you want to better understand what is going on during a
compilation, you can create a custom compiler, which is referred to as
backend in this section, that will print pretty print the fx
<code class="docutils literal notranslate"><span class="pre">GraphModule</span></code> extracted from Dynamo’s bytecode analysis
and return a <code class="docutils literal notranslate"><span class="pre">forward()</span></code> callable.</p>
<p>For example:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="k">def</span> <span class="nf">my_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">:</span> <span class="n">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="nb">print</span><span class="p">(</span><span class="s2">&quot;my_compiler() called with FX graph:&quot;</span><span class="p">)</span>
    <span class="n">gm</span><span class="o">.</span><span class="n">graph</span><span class="o">.</span><span class="n">print_tabular</span><span class="p">()</span>
    <span class="k">return</span> <span class="n">gm</span><span class="o">.</span><span class="n">forward</span>  <span class="c1"># return a python callable</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">backend</span><span class="o">=</span><span class="n">my_compiler</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">fn</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
    <span class="n">a</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cos</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
    <span class="n">b</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sin</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">a</span> <span class="o">+</span> <span class="n">b</span>
<span class="n">fn</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">))</span>
</pre></div>
</div>
<p>Running the above example produces the following output:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">my_compiler</span><span class="p">()</span> <span class="n">called</span> <span class="k">with</span> <span class="n">FX</span> <span class="n">graph</span><span class="p">:</span>
<span class="n">opcode</span>         <span class="n">name</span>    <span class="n">target</span>                                                  <span class="n">args</span>        <span class="n">kwargs</span>
<span class="o">-------------</span>  <span class="o">------</span>  <span class="o">------------------------------------------------------</span>  <span class="o">----------</span>  <span class="o">--------</span>
<span class="n">placeholder</span>    <span class="n">x</span>       <span class="n">x</span>                                                       <span class="p">()</span>          <span class="p">{}</span>
<span class="n">placeholder</span>    <span class="n">y</span>       <span class="n">y</span>                                                       <span class="p">()</span>          <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">cos</span>     <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">method</span> <span class="n">cos</span> <span class="n">of</span> <span class="nb">type</span> <span class="nb">object</span> <span class="n">at</span> <span class="mh">0x7f1a894649a8</span><span class="o">&gt;</span>  <span class="p">(</span><span class="n">x</span><span class="p">,)</span>        <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">sin</span>     <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">method</span> <span class="n">sin</span> <span class="n">of</span> <span class="nb">type</span> <span class="nb">object</span> <span class="n">at</span> <span class="mh">0x7f1a894649a8</span><span class="o">&gt;</span>  <span class="p">(</span><span class="n">y</span><span class="p">,)</span>        <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">add</span>     <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">add</span><span class="o">&gt;</span>                                 <span class="p">(</span><span class="n">cos</span><span class="p">,</span> <span class="n">sin</span><span class="p">)</span>  <span class="p">{}</span>
<span class="n">output</span>         <span class="n">output</span>  <span class="n">output</span>                                                  <span class="p">((</span><span class="n">add</span><span class="p">,),)</span>   <span class="p">{}</span>
</pre></div>
</div>
<p>This works for <code class="docutils literal notranslate"><span class="pre">torch.nn.Module</span></code> as well as shown below:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="k">def</span> <span class="nf">my_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">:</span> <span class="n">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="nb">print</span><span class="p">(</span><span class="s2">&quot;my_compiler() called with FX graph:&quot;</span><span class="p">)</span>
    <span class="n">gm</span><span class="o">.</span><span class="n">graph</span><span class="o">.</span><span class="n">print_tabular</span><span class="p">()</span>
    <span class="k">return</span> <span class="n">gm</span><span class="o">.</span><span class="n">forward</span>  <span class="c1"># return a python callable</span>
<span class="k">class</span> <span class="nc">MockModule</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">relu</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ReLU</span><span class="p">()</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">cos</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
<span class="n">mod</span> <span class="o">=</span> <span class="n">MockModule</span><span class="p">()</span>
<span class="n">optimized_mod</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">mod</span><span class="p">,</span> <span class="n">backend</span><span class="o">=</span><span class="n">my_compiler</span><span class="p">)</span>
<span class="n">optimized_mod</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">))</span>
</pre></div>
</div>
<p>Let’s take a look at one more example with control flow:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>
<span class="kn">import</span> <span class="nn">torch</span>
<span class="k">def</span> <span class="nf">my_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">:</span> <span class="n">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="nb">print</span><span class="p">(</span><span class="s2">&quot;my_compiler() called with FX graph:&quot;</span><span class="p">)</span>
    <span class="n">gm</span><span class="o">.</span><span class="n">graph</span><span class="o">.</span><span class="n">print_tabular</span><span class="p">()</span>
    <span class="k">return</span> <span class="n">gm</span><span class="o">.</span><span class="n">forward</span>  <span class="c1"># return a python callable</span>
<span class="nd">@torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">backend</span><span class="o">=</span><span class="n">my_compiler</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">toy_example</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">):</span>
    <span class="n">x</span> <span class="o">=</span> <span class="n">a</span> <span class="o">/</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">a</span><span class="p">)</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">b</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">b</span> <span class="o">=</span> <span class="n">b</span> <span class="o">*</span> <span class="o">-</span><span class="mi">1</span>
    <span class="k">return</span> <span class="n">x</span> <span class="o">*</span> <span class="n">b</span>
<span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">100</span><span class="p">):</span>
    <span class="n">toy_example</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">),</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">10</span><span class="p">))</span>
</pre></div>
</div>
<p>Running this example produces the following output:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">my_compiler</span><span class="p">()</span> <span class="n">called</span> <span class="k">with</span> <span class="n">FX</span> <span class="n">graph</span><span class="p">:</span>
<span class="n">opcode</span>         <span class="n">name</span>     <span class="n">target</span>                                                  <span class="n">args</span>              <span class="n">kwargs</span>
<span class="o">-------------</span>  <span class="o">-------</span>  <span class="o">------------------------------------------------------</span>  <span class="o">----------------</span>  <span class="o">--------</span>
<span class="n">placeholder</span>    <span class="n">a</span>        <span class="n">a</span>                                                       <span class="p">()</span>                <span class="p">{}</span>
<span class="n">placeholder</span>    <span class="n">b</span>        <span class="n">b</span>                                                       <span class="p">()</span>                <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">abs_1</span>    <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">method</span> <span class="nb">abs</span> <span class="n">of</span> <span class="nb">type</span> <span class="nb">object</span> <span class="n">at</span> <span class="mh">0x7f8d259298a0</span><span class="o">&gt;</span>  <span class="p">(</span><span class="n">a</span><span class="p">,)</span>              <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">add</span>      <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">add</span><span class="o">&gt;</span>                                 <span class="p">(</span><span class="n">abs_1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>        <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">truediv</span>  <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">truediv</span><span class="o">&gt;</span>                             <span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">add</span><span class="p">)</span>          <span class="p">{}</span>
<span class="n">call_method</span>    <span class="n">sum_1</span>    <span class="nb">sum</span>                                                     <span class="p">(</span><span class="n">b</span><span class="p">,)</span>              <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">lt</span>       <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">lt</span><span class="o">&gt;</span>                                  <span class="p">(</span><span class="n">sum_1</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>        <span class="p">{}</span>
<span class="n">output</span>         <span class="n">output</span>   <span class="n">output</span>                                                  <span class="p">((</span><span class="n">truediv</span><span class="p">,</span> <span class="n">lt</span><span class="p">),)</span>  <span class="p">{}</span>

<span class="n">my_compiler</span><span class="p">()</span> <span class="n">called</span> <span class="k">with</span> <span class="n">FX</span> <span class="n">graph</span><span class="p">:</span>
<span class="n">opcode</span>         <span class="n">name</span>    <span class="n">target</span>                   <span class="n">args</span>         <span class="n">kwargs</span>
<span class="o">-------------</span>  <span class="o">------</span>  <span class="o">-----------------------</span>  <span class="o">-----------</span>  <span class="o">--------</span>
<span class="n">placeholder</span>    <span class="n">b</span>       <span class="n">b</span>                        <span class="p">()</span>           <span class="p">{}</span>
<span class="n">placeholder</span>    <span class="n">x</span>       <span class="n">x</span>                        <span class="p">()</span>           <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">mul</span>     <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">mul</span><span class="o">&gt;</span>  <span class="p">(</span><span class="n">b</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>      <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">mul_1</span>   <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">mul</span><span class="o">&gt;</span>  <span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">mul</span><span class="p">)</span>     <span class="p">{}</span>
<span class="n">output</span>         <span class="n">output</span>  <span class="n">output</span>                   <span class="p">((</span><span class="n">mul_1</span><span class="p">,),)</span>  <span class="p">{}</span>

<span class="n">my_compiler</span><span class="p">()</span> <span class="n">called</span> <span class="k">with</span> <span class="n">FX</span> <span class="n">graph</span><span class="p">:</span>
<span class="n">opcode</span>         <span class="n">name</span>    <span class="n">target</span>                   <span class="n">args</span>       <span class="n">kwargs</span>
<span class="o">-------------</span>  <span class="o">------</span>  <span class="o">-----------------------</span>  <span class="o">---------</span>  <span class="o">--------</span>
<span class="n">placeholder</span>    <span class="n">b</span>       <span class="n">b</span>                        <span class="p">()</span>         <span class="p">{}</span>
<span class="n">placeholder</span>    <span class="n">x</span>       <span class="n">x</span>                        <span class="p">()</span>         <span class="p">{}</span>
<span class="n">call_function</span>  <span class="n">mul</span>     <span class="o">&lt;</span><span class="n">built</span><span class="o">-</span><span class="ow">in</span> <span class="n">function</span> <span class="n">mul</span><span class="o">&gt;</span>  <span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">b</span><span class="p">)</span>     <span class="p">{}</span>
<span class="n">output</span>         <span class="n">output</span>  <span class="n">output</span>                   <span class="p">((</span><span class="n">mul</span><span class="p">,),)</span>  <span class="p">{}</span>
</pre></div>
</div>
<p>The order of the last two graphs is nondeterministic depending
on which one is encountered first by the just-in-time compiler.</p>
</section>
<section id="speedy-backend">
<h3>Speedy Backend<a class="headerlink" href="#speedy-backend" title="Permalink to this heading">¶</a></h3>
<p>Integrating a custom backend that offers superior performance is also
easy and we’ll integrate a real one
with <a class="reference external" href="https://pytorch.org/docs/stable/generated/torch.jit.optimize_for_inference.html">optimize_for_inference</a>:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">optimize_for_inference_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">:</span> <span class="n">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="n">scripted</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">script</span><span class="p">(</span><span class="n">gm</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">optimize_for_inference</span><span class="p">(</span><span class="n">scripted</span><span class="p">)</span>
</pre></div>
</div>
<p>And then you should be able to optimize any existing code with:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="nd">@torch</span><span class="o">.</span><span class="n">compile</span><span class="p">(</span><span class="n">backend</span><span class="o">=</span><span class="n">optimize_for_inference_compiler</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">code_to_accelerate</span><span class="p">():</span>
    <span class="o">...</span>
</pre></div>
</div>
</section>
<section id="composable-backends">
<h3>Composable Backends<a class="headerlink" href="#composable-backends" title="Permalink to this heading">¶</a></h3>
<p>TorchDynamo includes many backends, which can be found in
<a class="reference external" href="https://github.com/pytorch/pytorch/blob/master/torch/_dynamo/optimizations/backends.py">backends.py</a>
or <code class="docutils literal notranslate"><span class="pre">torch._dynamo.list_backends()</span></code>. You can combine these backends
together with the following code:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">torch._dynamo.optimizations</span> <span class="kn">import</span> <span class="n">BACKENDS</span>
 <span class="k">def</span> <span class="nf">my_compiler</span><span class="p">(</span><span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">:</span> <span class="n">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="k">try</span><span class="p">:</span>
         <span class="n">trt_compiled</span> <span class="o">=</span> <span class="n">BACKENDS</span><span class="p">[</span><span class="s2">&quot;tensorrt&quot;</span><span class="p">](</span><span class="n">gm</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">)</span>
         <span class="k">if</span> <span class="n">trt_compiled</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">trt_compiled</span>
     <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
         <span class="k">pass</span>
     <span class="c1"># first backend failed, try something else...</span>
     <span class="k">try</span><span class="p">:</span>
         <span class="n">inductor_compiled</span> <span class="o">=</span> <span class="n">BACKENDS</span><span class="p">[</span><span class="s2">&quot;inductor&quot;</span><span class="p">](</span><span class="n">gm</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">)</span>
         <span class="k">if</span> <span class="n">inductor_compiled</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">inductor_compiled</span>
     <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
         <span class="k">pass</span>
     <span class="k">return</span> <span class="n">gm</span><span class="o">.</span><span class="n">forward</span>
</pre></div>
</div>
</section>
</section>
</section>


             </article>
             
            </div>
            <footer>
  
    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="deep-dive.html" class="btn btn-neutral float-right" title="TorchDynamo Deeper Dive" accesskey="n" rel="next">Next <img src="../_static/images/chevron-right-orange.svg" class="next-page"></a>
      
      
        <a href="guards-overview.html" class="btn btn-neutral" title="Guards Overview" accesskey="p" rel="prev"><img src="../_static/images/chevron-right-orange.svg" class="previous-page"> Previous</a>
      
    </div>
  

    <hr>

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

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

</footer>

          </div>
<script>

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

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

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              <ul>
<li><a class="reference internal" href="#">Custom Backends</a><ul>
<li><a class="reference internal" href="#overview">Overview</a></li>
<li><a class="reference internal" href="#registering-custom-backends">Registering Custom Backends</a></li>
<li><a class="reference internal" href="#custom-backends-after-aotautograd">Custom Backends after AOTAutograd</a></li>
<li><a class="reference internal" href="#examples">Examples</a><ul>
<li><a class="reference internal" href="#debugging-backend">Debugging Backend</a></li>
<li><a class="reference internal" href="#speedy-backend">Speedy Backend</a></li>
<li><a class="reference internal" href="#composable-backends">Composable Backends</a></li>
</ul>
</li>
</ul>
</li>
</ul>

            </div>
          </div>
        </div>
      </section>
    </div>

  
       <script type="text/javascript" id="documentation_options" data-url_root="../" src="../_static/documentation_options.js"></script>
         <script data-url_root="../" id="documentation_options" src="../_static/documentation_options.js"></script>
         <script src="../_static/jquery.js"></script>
         <script src="../_static/underscore.js"></script>
         <script src="../_static/_sphinx_javascript_frameworks_compat.js"></script>
         <script src="../_static/doctools.js"></script>
         <script src="../_static/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="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="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>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>
            
          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="resources-mobile-menu-title" class="active">
            Docs
          </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/audio/stable/index.html">torchaudio</a>
            </li>

            <li>
              <a href="https://pytorch.org/text/stable/index.html">torchtext</a>
            </li>

            <li>
              <a href="https://pytorch.org/vision/stable/index.html">torchvision</a>
            </li>

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

            <li>
              <a href="https://pytorch.org/data">TorchData</a>
            </li>

            <li>
              <a href="https://pytorch.org/torchrec">TorchRec</a>
            </li>

            <li>
              <a href="https://pytorch.org/serve/">TorchServe</a>
            </li>

            <li>
              <a href="https://pytorch.org/torchx/">TorchX</a>
            </li>

            <li>
              <a href="https://pytorch.org/xla">PyTorch on XLA Devices</a>
            </li>
          </ul>

          <li class="resources-mobile-menu-title">
            Resources
          </li>
            
           <ul class="resources-mobile-menu-items">

            <li>
              <a href="https://pytorch.org/features">About</a>
            </li>

            <li>
              <a href="https://pytorch.org/foundation">PyTorch Foundation</a>
            </li>

            <li>
              <a href="https://pytorch.org/#community-module">Community</a>
            </li>

            <li>
              <a href="https://pytorch.org/community-stories">Community Stories</a>
            </li>

            <li>
              <a href="https://pytorch.org/resources">Developer Resources</a>
            </li>

            <li>
              <a href="https://pytorch.org/events">Events</a>
            </li>

            <li>
              <a href="https://discuss.pytorch.org/">Forums</a>
            </li>

            <li>
              <a href="https://pytorch.org/hub">Models (Beta)</a>
            </li>
          </ul>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </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>