2.0/_modules/torch/nn/functional.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 name="robots" content="noindex">
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>torch.nn.functional &mdash; PyTorch 2.0 documentation</title>
  

    <link rel="canonical" href="https://pytorch.org/docs/stable/_modules/torch/nn/functional.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" />


  <!-- 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>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/index.html">TorchDynamo Overview</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/installation.html">Installing TorchDynamo</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/get-started.html">Getting Started</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/guards-overview.html">Guards Overview</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/custom-backends.html">Custom Backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/deep-dive.html">TorchDynamo Deeper Dive</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/troubleshooting.html">TorchDynamo Troubleshooting</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dynamo/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">Module code</a> &gt;</li>
        
          <li><a href="../../torch.html">torch</a> &gt;</li>
        
      <li>torch.nn.functional</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">

        
          <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.nn.functional</h1><div class="highlight"><pre>
<span></span><span class="sd">&quot;&quot;&quot;Functional interface&quot;&quot;&quot;</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Union</span>
<span class="kn">import</span> <span class="nn">math</span>
<span class="kn">import</span> <span class="nn">warnings</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">_VF</span>
<span class="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">sym_int</span> <span class="k">as</span> <span class="n">_sym_int</span>
<span class="kn">from</span> <span class="nn">torch._C</span> <span class="kn">import</span> <span class="n">_infer_size</span><span class="p">,</span> <span class="n">_add_docstr</span>
<span class="kn">from</span> <span class="nn">torch._torch_docs</span> <span class="kn">import</span> <span class="n">reproducibility_notes</span><span class="p">,</span> <span class="n">tf32_notes</span><span class="p">,</span> <span class="n">sparse_support_notes</span>
<span class="c1"># A workaround to support both TorchScript and MyPy:</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">TYPE_CHECKING</span>
<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="kn">from</span> <span class="nn">torch.types</span> <span class="kn">import</span> <span class="n">_dtype</span> <span class="k">as</span> <span class="n">DType</span>
<span class="k">else</span><span class="p">:</span>
    <span class="c1"># The JIT doesn&#39;t understand Union, nor torch.dtype here</span>
    <span class="n">DType</span> <span class="o">=</span> <span class="nb">int</span>

<span class="kn">from</span> <span class="nn">.._jit_internal</span> <span class="kn">import</span> <span class="n">boolean_dispatch</span><span class="p">,</span> <span class="n">_overload</span><span class="p">,</span> <span class="n">BroadcastingList1</span><span class="p">,</span> <span class="n">BroadcastingList2</span><span class="p">,</span> <span class="n">BroadcastingList3</span>
<span class="kn">from</span> <span class="nn">..overrides</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">has_torch_function</span><span class="p">,</span> <span class="n">has_torch_function_unary</span><span class="p">,</span> <span class="n">has_torch_function_variadic</span><span class="p">,</span>
    <span class="n">handle_torch_function</span><span class="p">)</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">_reduction</span> <span class="k">as</span> <span class="n">_Reduction</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">grad</span>  <span class="c1"># noqa: F401</span>
<span class="kn">from</span> <span class="nn">.modules</span> <span class="kn">import</span> <span class="n">utils</span>
<span class="kn">from</span> <span class="nn">.modules.utils</span> <span class="kn">import</span> <span class="n">_single</span><span class="p">,</span> <span class="n">_pair</span><span class="p">,</span> <span class="n">_triple</span><span class="p">,</span> <span class="n">_list_with_default</span>

<span class="n">Tensor</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span>

<span class="n">conv1d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv1d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -&gt; Tensor</span>

<span class="sd">Applies a 1D convolution over an input signal composed of several input</span>
<span class="sd">planes.</span>

<span class="sd">{tf32_note}</span>

<span class="sd">See :class:`~torch.nn.Conv1d` for details and output shape.</span>

<span class="sd">Note:</span>
<span class="sd">    {cudnn_reproducibility_note}</span>

<span class="sd">Note:</span>
<span class="sd">    This operator supports complex data types i.e. ``complex32, complex64, complex128``.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">,</span> <span class="o">**</span><span class="n">tf32_notes</span>
    <span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    input: input tensor of shape :math:`(\text</span><span class="si">{minibatch}</span><span class="s2"> , \text{in\_channels} , iW)`</span>
<span class="s2">    weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text</span><span class="si">{groups}</span><span class="s2">} , kW)`</span>
<span class="s2">    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: ``None``</span>
<span class="s2">    stride: the stride of the convolving kernel. Can be a single number or</span>
<span class="s2">      a one-element tuple `(sW,)`. Default: 1</span>
<span class="s2">    padding: implicit paddings on both sides of the input. Can be a string {&#39;valid&#39;, &#39;same&#39;},</span>
<span class="s2">      single number or a one-element tuple `(padW,)`. Default: 0</span>
<span class="s2">      ``padding=&#39;valid&#39;`` is the same as no padding. ``padding=&#39;same&#39;`` pads</span>
<span class="s2">      the input so the output has the same shape as the input. However, this mode</span>
<span class="s2">      doesn&#39;t support any stride values other than 1.</span>

<span class="s2">      .. warning::</span>
<span class="s2">          For ``padding=&#39;same&#39;``, if the ``weight`` is even-length and</span>
<span class="s2">          ``dilation`` is odd in any dimension, a full :func:`pad` operation</span>
<span class="s2">          may be needed internally. Lowering performance.</span>
<span class="s2">    dilation: the spacing between kernel elements. Can be a single number or</span>
<span class="s2">      a one-element tuple `(dW,)`. Default: 1</span>
<span class="s2">    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by</span>
<span class="s2">      the number of groups. Default: 1</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; inputs = torch.randn(33, 16, 30)</span>
<span class="s2">    &gt;&gt;&gt; filters = torch.randn(20, 16, 5)</span>
<span class="s2">    &gt;&gt;&gt; F.conv1d(inputs, filters)</span>
<span class="s2">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">conv2d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv2d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -&gt; Tensor</span>

<span class="sd">Applies a 2D convolution over an input image composed of several input</span>
<span class="sd">planes.</span>

<span class="sd">{tf32_note}</span>

<span class="sd">See :class:`~torch.nn.Conv2d` for details and output shape.</span>

<span class="sd">Note:</span>
<span class="sd">    {cudnn_reproducibility_note}</span>

<span class="sd">Note:</span>
<span class="sd">    This operator supports complex data types i.e. ``complex32, complex64, complex128``.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">,</span> <span class="o">**</span><span class="n">tf32_notes</span>
    <span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    input: input tensor of shape :math:`(\text</span><span class="si">{minibatch}</span><span class="s2"> , \text{in\_channels} , iH , iW)`</span>
<span class="s2">    weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text</span><span class="si">{groups}</span><span class="s2">} , kH , kW)`</span>
<span class="s2">    bias: optional bias tensor of shape :math:`(\text{out\_channels})`. Default: ``None``</span>
<span class="s2">    stride: the stride of the convolving kernel. Can be a single number or a</span>
<span class="s2">      tuple `(sH, sW)`. Default: 1</span>
<span class="s2">    padding: implicit paddings on both sides of the input. Can be a string {&#39;valid&#39;, &#39;same&#39;},</span>
<span class="s2">      single number or a tuple `(padH, padW)`. Default: 0</span>
<span class="s2">      ``padding=&#39;valid&#39;`` is the same as no padding. ``padding=&#39;same&#39;`` pads</span>
<span class="s2">      the input so the output has the same shape as the input. However, this mode</span>
<span class="s2">      doesn&#39;t support any stride values other than 1.</span>

<span class="s2">      .. warning::</span>
<span class="s2">          For ``padding=&#39;same&#39;``, if the ``weight`` is even-length and</span>
<span class="s2">          ``dilation`` is odd in any dimension, a full :func:`pad` operation</span>
<span class="s2">          may be needed internally. Lowering performance.</span>

<span class="s2">    dilation: the spacing between kernel elements. Can be a single number or</span>
<span class="s2">      a tuple `(dH, dW)`. Default: 1</span>
<span class="s2">    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the</span>
<span class="s2">      number of groups. Default: 1</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; # With square kernels and equal stride</span>
<span class="s2">    &gt;&gt;&gt; filters = torch.randn(8, 4, 3, 3)</span>
<span class="s2">    &gt;&gt;&gt; inputs = torch.randn(1, 4, 5, 5)</span>
<span class="s2">    &gt;&gt;&gt; F.conv2d(inputs, filters, padding=1)</span>
<span class="s2">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>  <span class="c1"># noqa: E501</span>

<span class="n">conv3d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv3d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -&gt; Tensor</span>

<span class="sd">Applies a 3D convolution over an input image composed of several input</span>
<span class="sd">planes.</span>

<span class="sd">{tf32_note}</span>

<span class="sd">See :class:`~torch.nn.Conv3d` for details and output shape.</span>

<span class="sd">Note:</span>
<span class="sd">    {cudnn_reproducibility_note}</span>

<span class="sd">Note:</span>
<span class="sd">    This operator supports complex data types i.e. ``complex32, complex64, complex128``.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">,</span> <span class="o">**</span><span class="n">tf32_notes</span>
    <span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    input: input tensor of shape :math:`(\text</span><span class="si">{minibatch}</span><span class="s2"> , \text{in\_channels} , iT , iH , iW)`</span>
<span class="s2">    weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text</span><span class="si">{groups}</span><span class="s2">} , kT , kH , kW)`</span>
<span class="s2">    bias: optional bias tensor of shape :math:`(\text{out\_channels})`. Default: None</span>
<span class="s2">    stride: the stride of the convolving kernel. Can be a single number or a</span>
<span class="s2">      tuple `(sT, sH, sW)`. Default: 1</span>
<span class="s2">    padding: implicit paddings on both sides of the input. Can be a string {&#39;valid&#39;, &#39;same&#39;},</span>
<span class="s2">      single number or a tuple `(padT, padH, padW)`. Default: 0</span>
<span class="s2">      ``padding=&#39;valid&#39;`` is the same as no padding. ``padding=&#39;same&#39;`` pads</span>
<span class="s2">      the input so the output has the same shape as the input. However, this mode</span>
<span class="s2">      doesn&#39;t support any stride values other than 1.</span>

<span class="s2">      .. warning::</span>
<span class="s2">          For ``padding=&#39;same&#39;``, if the ``weight`` is even-length and</span>
<span class="s2">          ``dilation`` is odd in any dimension, a full :func:`pad` operation</span>
<span class="s2">          may be needed internally. Lowering performance.</span>

<span class="s2">    dilation: the spacing between kernel elements. Can be a single number or</span>
<span class="s2">      a tuple `(dT, dH, dW)`. Default: 1</span>
<span class="s2">    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by</span>
<span class="s2">      the number of groups. Default: 1</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; filters = torch.randn(33, 16, 3, 3, 3)</span>
<span class="s2">    &gt;&gt;&gt; inputs = torch.randn(20, 16, 50, 10, 20)</span>
<span class="s2">    &gt;&gt;&gt; F.conv3d(inputs, filters)</span>
<span class="s2">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>  <span class="c1"># noqa: E501</span>

<span class="n">conv_transpose1d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv_transpose1d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">conv_transpose1d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -&gt; Tensor</span>

<span class="sd">Applies a 1D transposed convolution operator over an input signal</span>
<span class="sd">composed of several input planes, sometimes also called &quot;deconvolution&quot;.</span>

<span class="sd">{tf32_note}</span>

<span class="sd">See :class:`~torch.nn.ConvTranspose1d` for details and output shape.</span>

<span class="sd">Note:</span>
<span class="sd">    {cudnn_reproducibility_note}</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">,</span> <span class="o">**</span><span class="n">tf32_notes</span>
    <span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    input: input tensor of shape :math:`(\text</span><span class="si">{minibatch}</span><span class="s2"> , \text{in\_channels} , iW)`</span>
<span class="s2">    weight: filters of shape :math:`(\text{in\_channels} , \frac{\text{out\_channels}}{\text</span><span class="si">{groups}</span><span class="s2">} , kW)`</span>
<span class="s2">    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: None</span>
<span class="s2">    stride: the stride of the convolving kernel. Can be a single number or a</span>
<span class="s2">      tuple ``(sW,)``. Default: 1</span>
<span class="s2">    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both</span>
<span class="s2">      sides of each dimension in the input. Can be a single number or a tuple</span>
<span class="s2">      ``(padW,)``. Default: 0</span>
<span class="s2">    output_padding: additional size added to one side of each dimension in the</span>
<span class="s2">      output shape. Can be a single number or a tuple ``(out_padW)``. Default: 0</span>
<span class="s2">    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the</span>
<span class="s2">      number of groups. Default: 1</span>
<span class="s2">    dilation: the spacing between kernel elements. Can be a single number or</span>
<span class="s2">      a tuple ``(dW,)``. Default: 1</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; inputs = torch.randn(20, 16, 50)</span>
<span class="s2">    &gt;&gt;&gt; weights = torch.randn(16, 33, 5)</span>
<span class="s2">    &gt;&gt;&gt; F.conv_transpose1d(inputs, weights)</span>
<span class="s2">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">conv_transpose2d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv_transpose2d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -&gt; Tensor</span>

<span class="sd">Applies a 2D transposed convolution operator over an input image</span>
<span class="sd">composed of several input planes, sometimes also called &quot;deconvolution&quot;.</span>

<span class="sd">{tf32_note}</span>

<span class="sd">See :class:`~torch.nn.ConvTranspose2d` for details and output shape.</span>

<span class="sd">Note:</span>
<span class="sd">    {cudnn_reproducibility_note}</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">,</span> <span class="o">**</span><span class="n">tf32_notes</span>
    <span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    input: input tensor of shape :math:`(\text</span><span class="si">{minibatch}</span><span class="s2"> , \text{in\_channels} , iH , iW)`</span>
<span class="s2">    weight: filters of shape :math:`(\text{in\_channels} , \frac{\text{out\_channels}}{\text</span><span class="si">{groups}</span><span class="s2">} , kH , kW)`</span>
<span class="s2">    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: None</span>
<span class="s2">    stride: the stride of the convolving kernel. Can be a single number or a</span>
<span class="s2">      tuple ``(sH, sW)``. Default: 1</span>
<span class="s2">    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both</span>
<span class="s2">      sides of each dimension in the input. Can be a single number or a tuple</span>
<span class="s2">      ``(padH, padW)``. Default: 0</span>
<span class="s2">    output_padding: additional size added to one side of each dimension in the</span>
<span class="s2">      output shape. Can be a single number or a tuple ``(out_padH, out_padW)``.</span>
<span class="s2">      Default: 0</span>
<span class="s2">    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the</span>
<span class="s2">      number of groups. Default: 1</span>
<span class="s2">    dilation: the spacing between kernel elements. Can be a single number or</span>
<span class="s2">      a tuple ``(dH, dW)``. Default: 1</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; # With square kernels and equal stride</span>
<span class="s2">    &gt;&gt;&gt; inputs = torch.randn(1, 4, 5, 5)</span>
<span class="s2">    &gt;&gt;&gt; weights = torch.randn(4, 8, 3, 3)</span>
<span class="s2">    &gt;&gt;&gt; F.conv_transpose2d(inputs, weights, padding=1)</span>
<span class="s2">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>  <span class="c1"># noqa: E501</span>

<span class="n">conv_transpose3d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv_transpose3d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">conv_transpose3d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -&gt; Tensor</span>

<span class="sd">Applies a 3D transposed convolution operator over an input image</span>
<span class="sd">composed of several input planes, sometimes also called &quot;deconvolution&quot;</span>

<span class="sd">{tf32_note}</span>

<span class="sd">See :class:`~torch.nn.ConvTranspose3d` for details and output shape.</span>

<span class="sd">Note:</span>
<span class="sd">    {cudnn_reproducibility_note}</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
        <span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">,</span> <span class="o">**</span><span class="n">tf32_notes</span>
    <span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    input: input tensor of shape :math:`(\text</span><span class="si">{minibatch}</span><span class="s2"> , \text{in\_channels} , iT , iH , iW)`</span>
<span class="s2">    weight: filters of shape :math:`(\text{in\_channels} , \frac{\text{out\_channels}}{\text</span><span class="si">{groups}</span><span class="s2">} , kT , kH , kW)`</span>
<span class="s2">    bias: optional bias of shape :math:`(\text{out\_channels})`. Default: None</span>
<span class="s2">    stride: the stride of the convolving kernel. Can be a single number or a</span>
<span class="s2">      tuple ``(sT, sH, sW)``. Default: 1</span>
<span class="s2">    padding: ``dilation * (kernel_size - 1) - padding`` zero-padding will be added to both</span>
<span class="s2">      sides of each dimension in the input. Can be a single number or a tuple</span>
<span class="s2">      ``(padT, padH, padW)``. Default: 0</span>
<span class="s2">    output_padding: additional size added to one side of each dimension in the</span>
<span class="s2">      output shape. Can be a single number or a tuple</span>
<span class="s2">      ``(out_padT, out_padH, out_padW)``. Default: 0</span>
<span class="s2">    groups: split input into groups, :math:`\text{in\_channels}` should be divisible by the</span>
<span class="s2">      number of groups. Default: 1</span>
<span class="s2">    dilation: the spacing between kernel elements. Can be a single number or</span>
<span class="s2">      a tuple `(dT, dH, dW)`. Default: 1</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; inputs = torch.randn(20, 16, 50, 10, 20)</span>
<span class="s2">    &gt;&gt;&gt; weights = torch.randn(16, 33, 3, 3, 3)</span>
<span class="s2">    &gt;&gt;&gt; F.conv_transpose3d(inputs, weights)</span>
<span class="s2">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>  <span class="c1"># noqa: E501</span>

<span class="n">conv_tbc</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">conv_tbc</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">Applies a 1-dimensional sequence convolution over an input sequence.</span>
<span class="sd">Input and output dimensions are (Time, Batch, Channels) - hence TBC.</span>

<span class="sd">Args:</span>
<span class="sd">    input: input tensor of shape :math:`(\text{sequence length} \times batch \times \text{in\_channels})`</span>
<span class="sd">    weight: filter of shape (:math:`\text{kernel width} \times \text{in\_channels} \times \text{out\_channels}`)</span>
<span class="sd">    bias: bias of shape (:math:`\text{out\_channels}`)</span>
<span class="sd">    pad: number of timesteps to pad. Default: 0</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="c1"># Pooling</span>
<span class="n">avg_pool1d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">avg_pool1d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) -&gt; Tensor</span>

<span class="sd">Applies a 1D average pooling over an input signal composed of several</span>
<span class="sd">input planes.</span>

<span class="sd">See :class:`~torch.nn.AvgPool1d` for details and output shape.</span>

<span class="sd">Args:</span>
<span class="sd">    input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`</span>
<span class="sd">    kernel_size: the size of the window. Can be a single number or a</span>
<span class="sd">      tuple `(kW,)`</span>
<span class="sd">    stride: the stride of the window. Can be a single number or a tuple</span>
<span class="sd">      `(sW,)`. Default: :attr:`kernel_size`</span>
<span class="sd">    padding: implicit zero paddings on both sides of the input. Can be a</span>
<span class="sd">      single number or a tuple `(padW,)`. Default: 0</span>
<span class="sd">    ceil_mode: when True, will use `ceil` instead of `floor` to compute the</span>
<span class="sd">        output shape. Default: ``False``</span>
<span class="sd">    count_include_pad: when True, will include the zero-padding in the</span>
<span class="sd">        averaging calculation. Default: ``True``</span>

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

<span class="sd">    &gt;&gt;&gt; # pool of square window of size=3, stride=2</span>
<span class="sd">    &gt;&gt;&gt; input = torch.tensor([[[1, 2, 3, 4, 5, 6, 7]]], dtype=torch.float32)</span>
<span class="sd">    &gt;&gt;&gt; F.avg_pool1d(input, kernel_size=3, stride=2)</span>
<span class="sd">    tensor([[[ 2.,  4.,  6.]]])</span>

<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="n">avg_pool2d</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">avg_pool2d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -&gt; Tensor</span>

<span class="sd">Applies 2D average-pooling operation in :math:`kH \times kW` regions by step size</span>
<span class="sd">:math:`sH \times sW` steps. The number of output features is equal to the number of</span>
<span class="sd">input planes.</span>

<span class="sd">See :class:`~torch.nn.AvgPool2d` for details and output shape.</span>

<span class="sd">Args:</span>
<span class="sd">    input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`</span>
<span class="sd">    kernel_size: size of the pooling region. Can be a single number or a</span>
<span class="sd">      tuple `(kH, kW)`</span>
<span class="sd">    stride: stride of the pooling operation. Can be a single number or a</span>
<span class="sd">      tuple `(sH, sW)`. Default: :attr:`kernel_size`</span>
<span class="sd">    padding: implicit zero paddings on both sides of the input. Can be a</span>
<span class="sd">      single number or a tuple `(padH, padW)`. Default: 0</span>
<span class="sd">    ceil_mode: when True, will use `ceil` instead of `floor` in the formula</span>
<span class="sd">        to compute the output shape. Default: ``False``</span>
<span class="sd">    count_include_pad: when True, will include the zero-padding in the</span>
<span class="sd">        averaging calculation. Default: ``True``</span>
<span class="sd">    divisor_override: if specified, it will be used as divisor, otherwise</span>
<span class="sd">         size of the pooling region will be used. Default: None</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">avg_pool3d</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">avg_pool3d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -&gt; Tensor</span>

<span class="sd">Applies 3D average-pooling operation in :math:`kT \times kH \times kW` regions by step</span>
<span class="sd">size :math:`sT \times sH \times sW` steps. The number of output features is equal to</span>
<span class="sd">:math:`\lfloor\frac{\text{input planes}}{sT}\rfloor`.</span>

<span class="sd">See :class:`~torch.nn.AvgPool3d` for details and output shape.</span>

<span class="sd">Args:</span>
<span class="sd">    input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iT \times iH , iW)`</span>
<span class="sd">    kernel_size: size of the pooling region. Can be a single number or a</span>
<span class="sd">      tuple `(kT, kH, kW)`</span>
<span class="sd">    stride: stride of the pooling operation. Can be a single number or a</span>
<span class="sd">      tuple `(sT, sH, sW)`. Default: :attr:`kernel_size`</span>
<span class="sd">    padding: implicit zero paddings on both sides of the input. Can be a</span>
<span class="sd">      single number or a tuple `(padT, padH, padW)`, Default: 0</span>
<span class="sd">    ceil_mode: when True, will use `ceil` instead of `floor` in the formula</span>
<span class="sd">        to compute the output shape</span>
<span class="sd">    count_include_pad: when True, will include the zero-padding in the</span>
<span class="sd">        averaging calculation</span>
<span class="sd">    divisor_override: if specified, it will be used as divisor, otherwise</span>
<span class="sd">        size of the pooling region will be used. Default: None</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">fractional_max_pool2d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">,</span>
    <span class="n">output_ratio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">return_indices</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">_random_samples</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies 2D fractional max pooling over an input signal composed of several input planes.</span>

<span class="sd">    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham</span>

<span class="sd">    The max-pooling operation is applied in :math:`kH \times kW` regions by a stochastic</span>
<span class="sd">    step size determined by the target output size.</span>
<span class="sd">    The number of output features is equal to the number of input planes.</span>

<span class="sd">    Args:</span>
<span class="sd">        kernel_size: the size of the window to take a max over.</span>
<span class="sd">                     Can be a single number :math:`k` (for a square kernel of :math:`k \times k`)</span>
<span class="sd">                     or a tuple `(kH, kW)`</span>
<span class="sd">        output_size: the target output size of the image of the form :math:`oH \times oW`.</span>
<span class="sd">                     Can be a tuple `(oH, oW)` or a single number :math:`oH` for a square image :math:`oH \times oH`</span>
<span class="sd">        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.</span>
<span class="sd">                      This has to be a number or tuple in the range (0, 1)</span>
<span class="sd">        return_indices: if ``True``, will return the indices along with the outputs.</span>
<span class="sd">                        Useful to pass to :func:`~torch.nn.functional.max_unpool2d`.</span>

<span class="sd">    Examples::</span>
<span class="sd">        &gt;&gt;&gt; input = torch.randn(20, 16, 50, 32)</span>
<span class="sd">        &gt;&gt;&gt; # pool of square window of size=3, and target output size 13x12</span>
<span class="sd">        &gt;&gt;&gt; F.fractional_max_pool2d(input, 3, output_size=(13, 12))</span>
<span class="sd">        &gt;&gt;&gt; # pool of square window and target output size being half of input image size</span>
<span class="sd">        &gt;&gt;&gt; F.fractional_max_pool2d(input, 3, output_ratio=(0.5, 0.5))</span>

<span class="sd">    .. _Fractional MaxPooling:</span>
<span class="sd">        http://arxiv.org/abs/1412.6071</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">fractional_max_pool2d_with_indices</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
            <span class="n">output_ratio</span><span class="o">=</span><span class="n">output_ratio</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
            <span class="n">_random_samples</span><span class="o">=</span><span class="n">_random_samples</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">output_ratio</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;fractional_max_pool2d requires specifying either &quot;</span> <span class="s2">&quot;an output_size or an output_ratio&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">output_ratio</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="n">_output_ratio</span> <span class="o">=</span> <span class="n">_pair</span><span class="p">(</span><span class="n">output_ratio</span><span class="p">)</span>
        <span class="n">output_size</span> <span class="o">=</span> <span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">)</span> <span class="o">*</span> <span class="n">_output_ratio</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span> <span class="nb">int</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="n">_output_ratio</span><span class="p">[</span><span class="mi">1</span><span class="p">])]</span>

    <span class="k">if</span> <span class="n">_random_samples</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">n_batch</span> <span class="o">=</span> <span class="mi">1</span> <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="k">else</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
        <span class="n">_random_samples</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">n_batch</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">3</span><span class="p">),</span> <span class="mi">2</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="nb">input</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="nb">input</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">fractional_max_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_fractional_max_pool2d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">,</span>
    <span class="n">output_ratio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">return_indices</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">_random_samples</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">fractional_max_pool2d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
            <span class="n">output_ratio</span><span class="o">=</span><span class="n">output_ratio</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
            <span class="n">_random_samples</span><span class="o">=</span><span class="n">_random_samples</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">fractional_max_pool2d_with_indices</span><span class="p">(</span>
        <span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">output_ratio</span><span class="p">,</span> <span class="n">return_indices</span><span class="p">,</span> <span class="n">_random_samples</span>
    <span class="p">)[</span><span class="mi">0</span><span class="p">]</span>


<span class="n">fractional_max_pool2d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">fractional_max_pool2d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_fractional_max_pool2d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;fractional_max_pool2d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">fractional_max_pool3d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</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="p">,</span>
    <span class="n">output_ratio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">return_indices</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">_random_samples</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies 3D fractional max pooling over an input signal composed of several input planes.</span>

<span class="sd">    Fractional MaxPooling is described in detail in the paper `Fractional MaxPooling`_ by Ben Graham</span>

<span class="sd">    The max-pooling operation is applied in :math:`kT \times kH \times kW` regions by a stochastic</span>
<span class="sd">    step size determined by the target output size.</span>
<span class="sd">    The number of output features is equal to the number of input planes.</span>

<span class="sd">    Args:</span>
<span class="sd">        kernel_size: the size of the window to take a max over.</span>
<span class="sd">                     Can be a single number :math:`k` (for a square kernel of :math:`k \times k \times k`)</span>
<span class="sd">                     or a tuple `(kT, kH, kW)`</span>
<span class="sd">        output_size: the target output size of the form :math:`oT \times oH \times oW`.</span>
<span class="sd">                     Can be a tuple `(oT, oH, oW)` or a single number :math:`oH` for a cubic output</span>
<span class="sd">                     :math:`oH \times oH \times oH`</span>
<span class="sd">        output_ratio: If one wants to have an output size as a ratio of the input size, this option can be given.</span>
<span class="sd">                      This has to be a number or tuple in the range (0, 1)</span>
<span class="sd">        return_indices: if ``True``, will return the indices along with the outputs.</span>
<span class="sd">                        Useful to pass to :func:`~torch.nn.functional.max_unpool3d`.</span>

<span class="sd">    Shape:</span>
<span class="sd">        - Input: :math:`(N, C, T_{in}, H_{in}, W_{in})` or :math:`(C, T_{in}, H_{in}, W_{in})`.</span>
<span class="sd">        - Output: :math:`(N, C, T_{out}, H_{out}, W_{out})` or :math:`(C, T_{out}, H_{out}, W_{out})`, where</span>
<span class="sd">          :math:`(T_{out}, H_{out}, W_{out})=\text{output\_size}` or</span>
<span class="sd">          :math:`(T_{out}, H_{out}, W_{out})=\text{output\_ratio} \times (T_{in}, H_{in}, W_{in})`</span>

<span class="sd">    Examples::</span>
<span class="sd">        &gt;&gt;&gt; input = torch.randn(20, 16, 50, 32, 16)</span>
<span class="sd">        &gt;&gt;&gt; # pool of cubic window of size=3, and target output size 13x12x11</span>
<span class="sd">        &gt;&gt;&gt; F.fractional_max_pool3d(input, 3, output_size=(13, 12, 11))</span>
<span class="sd">        &gt;&gt;&gt; # pool of cubic window and target output size being half of input size</span>
<span class="sd">        &gt;&gt;&gt; F.fractional_max_pool3d(input, 3, output_ratio=(0.5, 0.5, 0.5))</span>

<span class="sd">    .. _Fractional MaxPooling:</span>
<span class="sd">        http://arxiv.org/abs/1412.6071</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">fractional_max_pool3d_with_indices</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
            <span class="n">output_ratio</span><span class="o">=</span><span class="n">output_ratio</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
            <span class="n">_random_samples</span><span class="o">=</span><span class="n">_random_samples</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">output_ratio</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;fractional_max_pool3d requires specifying either &quot;</span> <span class="s2">&quot;an output_size or an output_ratio&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">output_ratio</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="n">_output_ratio</span> <span class="o">=</span> <span class="n">_triple</span><span class="p">(</span><span class="n">output_ratio</span><span class="p">)</span>
        <span class="n">output_size</span> <span class="o">=</span> <span class="p">[</span>
            <span class="nb">int</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">3</span><span class="p">)</span> <span class="o">*</span> <span class="n">_output_ratio</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span>
            <span class="nb">int</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">)</span> <span class="o">*</span> <span class="n">_output_ratio</span><span class="p">[</span><span class="mi">1</span><span class="p">]),</span>
            <span class="nb">int</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="n">_output_ratio</span><span class="p">[</span><span class="mi">2</span><span class="p">]),</span>
        <span class="p">]</span>

    <span class="k">if</span> <span class="n">_random_samples</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">n_batch</span> <span class="o">=</span> <span class="mi">1</span> <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="k">else</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
        <span class="n">_random_samples</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">n_batch</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">4</span><span class="p">),</span> <span class="mi">3</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="nb">input</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="nb">input</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">fractional_max_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_fractional_max_pool3d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</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="p">,</span>
    <span class="n">output_ratio</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">return_indices</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">_random_samples</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">fractional_max_pool3d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_random_samples</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
            <span class="n">output_ratio</span><span class="o">=</span><span class="n">output_ratio</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
            <span class="n">_random_samples</span><span class="o">=</span><span class="n">_random_samples</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">fractional_max_pool3d_with_indices</span><span class="p">(</span>
        <span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">output_ratio</span><span class="p">,</span> <span class="n">return_indices</span><span class="p">,</span> <span class="n">_random_samples</span>
    <span class="p">)[</span><span class="mi">0</span><span class="p">]</span>


<span class="n">fractional_max_pool3d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">fractional_max_pool3d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_fractional_max_pool3d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;fractional_max_pool3d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">max_pool1d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList1</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">ceil_mode</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">return_indices</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="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)</span>

<span class="sd">    Applies a 1D max pooling over an input signal composed of several input</span>
<span class="sd">    planes.</span>

<span class="sd">    .. note::</span>
<span class="sd">        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from</span>
<span class="sd">        what seen in :class:`~torch.nn.MaxPool1d`, and will change in a future release.</span>

<span class="sd">    See :class:`~torch.nn.MaxPool1d` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`, minibatch dim optional.</span>
<span class="sd">        kernel_size: the size of the window. Can be a single number or a</span>
<span class="sd">            tuple `(kW,)`</span>
<span class="sd">        stride: the stride of the window. Can be a single number or a tuple</span>
<span class="sd">            `(sW,)`. Default: :attr:`kernel_size`</span>
<span class="sd">        padding: Implicit negative infinity padding to be added on both sides, must be &gt;= 0 and &lt;= kernel_size / 2.</span>
<span class="sd">        dilation: The stride between elements within a sliding window, must be &gt; 0.</span>
<span class="sd">        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This</span>
<span class="sd">                   ensures that every element in the input tensor is covered by a sliding window.</span>
<span class="sd">        return_indices: If ``True``, will return the argmax along with the max values.</span>
<span class="sd">                        Useful for :class:`torch.nn.functional.max_unpool1d` later</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_pool1d_with_indices</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
            <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">stride</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">max_pool1d_with_indices</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_max_pool1d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList1</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">ceil_mode</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">return_indices</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_pool1d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
            <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">stride</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">max_pool1d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>


<span class="n">max_pool1d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">max_pool1d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_max_pool1d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;max_pool1d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">max_pool2d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">ceil_mode</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">return_indices</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="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)</span>

<span class="sd">    Applies a 2D max pooling over an input signal composed of several input</span>
<span class="sd">    planes.</span>

<span class="sd">    .. note::</span>
<span class="sd">        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from</span>
<span class="sd">        what seen in :class:`~torch.nn.MaxPool2d`, and will change in a future release.</span>

<span class="sd">    See :class:`~torch.nn.MaxPool2d` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`, minibatch dim optional.</span>
<span class="sd">        kernel_size: size of the pooling region. Can be a single number or a</span>
<span class="sd">            tuple `(kH, kW)`</span>
<span class="sd">        stride: stride of the pooling operation. Can be a single number or a</span>
<span class="sd">            tuple `(sH, sW)`. Default: :attr:`kernel_size`</span>
<span class="sd">        padding: Implicit negative infinity padding to be added on both sides, must be &gt;= 0 and &lt;= kernel_size / 2.</span>
<span class="sd">        dilation: The stride between elements within a sliding window, must be &gt; 0.</span>
<span class="sd">        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This</span>
<span class="sd">                   ensures that every element in the input tensor is covered by a sliding window.</span>
<span class="sd">        return_indices: If ``True``, will return the argmax along with the max values.</span>
<span class="sd">                        Useful for :class:`torch.nn.functional.max_unpool2d` later</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_pool2d_with_indices</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
            <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">stride</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">max_pool2d_with_indices</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_max_pool2d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">ceil_mode</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">return_indices</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_pool2d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
            <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">stride</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">max_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>


<span class="n">max_pool2d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">max_pool2d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_max_pool2d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;max_pool2d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">max_pool3d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">ceil_mode</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">return_indices</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="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False)</span>

<span class="sd">    Applies a 3D max pooling over an input signal composed of several input</span>
<span class="sd">    planes.</span>

<span class="sd">    .. note::</span>
<span class="sd">        The order of :attr:`ceil_mode` and :attr:`return_indices` is different from</span>
<span class="sd">        what seen in :class:`~torch.nn.MaxPool3d`, and will change in a future release.</span>

<span class="sd">    See :class:`~torch.nn.MaxPool3d` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: input tensor :math:`(\text{minibatch} , \text{in\_channels} , iD, iH , iW)`, minibatch dim optional.</span>
<span class="sd">        kernel_size: size of the pooling region. Can be a single number or a</span>
<span class="sd">                     tuple `(kT, kH, kW)`</span>
<span class="sd">        stride: stride of the pooling operation. Can be a single number or a</span>
<span class="sd">                tuple `(sT, sH, sW)`. Default: :attr:`kernel_size`</span>
<span class="sd">        padding: Implicit negative infinity padding to be added on both sides, must be &gt;= 0 and &lt;= kernel_size / 2.</span>
<span class="sd">        dilation: The stride between elements within a sliding window, must be &gt; 0.</span>
<span class="sd">        ceil_mode: If ``True``, will use `ceil` instead of `floor` to compute the output shape. This</span>
<span class="sd">                   ensures that every element in the input tensor is covered by a sliding window.</span>
<span class="sd">        return_indices: If ``True``, will return the argmax along with the max values.</span>
<span class="sd">                        Useful for :class:`torch.nn.functional.max_unpool3d` later</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_pool3d_with_indices</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
            <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">stride</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">max_pool3d_with_indices</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_max_pool3d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">ceil_mode</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">return_indices</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_pool3d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span>
            <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">,</span>
            <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">stride</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">max_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">dilation</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>


<span class="n">max_pool3d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">max_pool3d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_max_pool3d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;max_pool3d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">_unpool_output_size</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">stride</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">padding</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]]</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">]:</span>
    <span class="n">input_size</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()</span>
    <span class="n">default_size</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">annotate</span><span class="p">(</span><span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="p">[])</span>
    <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)):</span>
        <span class="n">default_size</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">input_size</span><span class="p">[</span><span class="o">-</span><span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span> <span class="o">+</span> <span class="n">d</span><span class="p">]</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="n">stride</span><span class="p">[</span><span class="n">d</span><span class="p">]</span> <span class="o">+</span> <span class="n">kernel_size</span><span class="p">[</span><span class="n">d</span><span class="p">]</span> <span class="o">-</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">padding</span><span class="p">[</span><span class="n">d</span><span class="p">])</span>
    <span class="k">if</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">default_size</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">output_size</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span> <span class="o">+</span> <span class="mi">2</span><span class="p">:</span>
            <span class="n">output_size</span> <span class="o">=</span> <span class="n">output_size</span><span class="p">[</span><span class="mi">2</span><span class="p">:]</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">output_size</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;output_size should be a sequence containing &quot;</span>
                <span class="s2">&quot;</span><span class="si">{}</span><span class="s2"> or </span><span class="si">{}</span><span class="s2"> elements, but it has a length of &#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                    <span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span> <span class="o">+</span> <span class="mi">2</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">output_size</span><span class="p">)</span>
                <span class="p">)</span>
            <span class="p">)</span>
        <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)):</span>
            <span class="n">min_size</span> <span class="o">=</span> <span class="n">default_size</span><span class="p">[</span><span class="n">d</span><span class="p">]</span> <span class="o">-</span> <span class="n">stride</span><span class="p">[</span><span class="n">d</span><span class="p">]</span>
            <span class="n">max_size</span> <span class="o">=</span> <span class="n">default_size</span><span class="p">[</span><span class="n">d</span><span class="p">]</span> <span class="o">+</span> <span class="n">stride</span><span class="p">[</span><span class="n">d</span><span class="p">]</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">min_size</span> <span class="o">&lt;</span> <span class="n">output_size</span><span class="p">[</span><span class="n">d</span><span class="p">]</span> <span class="o">&lt;</span> <span class="n">max_size</span><span class="p">):</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s1">&#39;invalid output_size &quot;</span><span class="si">{}</span><span class="s1">&quot; (dim </span><span class="si">{}</span><span class="s1"> must be between </span><span class="si">{}</span><span class="s1"> and </span><span class="si">{}</span><span class="s1">)&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                        <span class="n">output_size</span><span class="p">,</span> <span class="n">d</span><span class="p">,</span> <span class="n">min_size</span><span class="p">,</span> <span class="n">max_size</span>
                    <span class="p">)</span>
                <span class="p">)</span>

        <span class="n">ret</span> <span class="o">=</span> <span class="n">output_size</span>
    <span class="k">return</span> <span class="n">ret</span>


<div class="viewcode-block" id="max_unpool1d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.max_unpool1d.html#torch.nn.functional.max_unpool1d">[docs]</a><span class="k">def</span> <span class="nf">max_unpool1d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">indices</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList1</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList1</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Computes a partial inverse of :class:`MaxPool1d`.</span>

<span class="sd">    See :class:`~torch.nn.MaxUnpool1d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_unpool1d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">indices</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="n">kernel_size</span> <span class="o">=</span> <span class="n">_single</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">_stride</span> <span class="o">=</span> <span class="n">_single</span><span class="p">(</span><span class="n">stride</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">_stride</span> <span class="o">=</span> <span class="n">kernel_size</span>
    <span class="n">padding</span> <span class="o">=</span> <span class="n">_single</span><span class="p">(</span><span class="n">padding</span><span class="p">)</span>
    <span class="n">output_size</span> <span class="o">=</span> <span class="n">_unpool_output_size</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">_stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">output_size</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
        <span class="n">output_size</span> <span class="o">=</span> <span class="n">output_size</span> <span class="o">+</span> <span class="p">[</span><span class="mi">1</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">output_size</span> <span class="o">=</span> <span class="n">output_size</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">max_unpool2d</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">indices</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">output_size</span><span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span></div>


<div class="viewcode-block" id="max_unpool2d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.max_unpool2d.html#torch.nn.functional.max_unpool2d">[docs]</a><span class="k">def</span> <span class="nf">max_unpool2d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">indices</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Computes a partial inverse of :class:`MaxPool2d`.</span>

<span class="sd">    See :class:`~torch.nn.MaxUnpool2d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_unpool2d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">indices</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="n">kernel_size</span> <span class="o">=</span> <span class="n">_pair</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">_stride</span> <span class="o">=</span> <span class="n">_pair</span><span class="p">(</span><span class="n">stride</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">_stride</span> <span class="o">=</span> <span class="n">kernel_size</span>
    <span class="n">padding</span> <span class="o">=</span> <span class="n">_pair</span><span class="p">(</span><span class="n">padding</span><span class="p">)</span>
    <span class="n">output_size</span> <span class="o">=</span> <span class="n">_unpool_output_size</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">_stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">max_unpool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span></div>


<div class="viewcode-block" id="max_unpool3d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.max_unpool3d.html#torch.nn.functional.max_unpool3d">[docs]</a><span class="k">def</span> <span class="nf">max_unpool3d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">indices</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</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="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">output_size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList3</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Computes a partial inverse of :class:`MaxPool3d`.</span>

<span class="sd">    See :class:`~torch.nn.MaxUnpool3d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">max_unpool3d</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">indices</span><span class="p">,</span>
            <span class="n">kernel_size</span><span class="p">,</span>
            <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span>
            <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span>
            <span class="n">output_size</span><span class="o">=</span><span class="n">output_size</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="n">kernel_size</span> <span class="o">=</span> <span class="n">_triple</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">_stride</span> <span class="o">=</span> <span class="n">_triple</span><span class="p">(</span><span class="n">stride</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">_stride</span> <span class="o">=</span> <span class="n">kernel_size</span>
    <span class="n">padding</span> <span class="o">=</span> <span class="n">_triple</span><span class="p">(</span><span class="n">padding</span><span class="p">)</span>
    <span class="n">output_size</span> <span class="o">=</span> <span class="n">_unpool_output_size</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">_stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">max_unpool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">_stride</span><span class="p">,</span> <span class="n">padding</span><span class="p">)</span></div>


<div class="viewcode-block" id="lp_pool2d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.lp_pool2d.html#torch.nn.functional.lp_pool2d">[docs]</a><span class="k">def</span> <span class="nf">lp_pool2d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">:</span> <span class="n">Union</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="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList2</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="p">,</span>
    <span class="n">ceil_mode</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a 2D power-average pooling over an input signal composed of</span>
<span class="sd">    several input planes. If the sum of all inputs to the power of `p` is</span>
<span class="sd">    zero, the gradient is set to zero as well.</span>

<span class="sd">    See :class:`~torch.nn.LPPool2d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">lp_pool2d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span>
        <span class="p">)</span>
    <span class="n">kw</span><span class="p">,</span> <span class="n">kh</span> <span class="o">=</span> <span class="n">utils</span><span class="o">.</span><span class="n">_pair</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">out</span> <span class="o">=</span> <span class="n">avg_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">norm_type</span><span class="p">),</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">out</span> <span class="o">=</span> <span class="n">avg_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">norm_type</span><span class="p">),</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">)</span>

    <span class="k">return</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">out</span><span class="p">)</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">abs</span><span class="p">(</span><span class="n">out</span><span class="p">)))</span><span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="n">kw</span> <span class="o">*</span> <span class="n">kh</span><span class="p">)</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">/</span> <span class="n">norm_type</span><span class="p">)</span></div>


<div class="viewcode-block" id="lp_pool1d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.lp_pool1d.html#torch.nn.functional.lp_pool1d">[docs]</a><span class="k">def</span> <span class="nf">lp_pool1d</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">:</span> <span class="n">Union</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="n">kernel_size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BroadcastingList1</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="p">,</span>
    <span class="n">ceil_mode</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a 1D power-average pooling over an input signal composed of</span>
<span class="sd">    several input planes. If the sum of all inputs to the power of `p` is</span>
<span class="sd">    zero, the gradient is set to zero as well.</span>

<span class="sd">    See :class:`~torch.nn.LPPool1d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">lp_pool1d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="o">=</span><span class="n">stride</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">stride</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">out</span> <span class="o">=</span> <span class="n">avg_pool1d</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">norm_type</span><span class="p">),</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">stride</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">out</span> <span class="o">=</span> <span class="n">avg_pool1d</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">norm_type</span><span class="p">),</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">ceil_mode</span><span class="o">=</span><span class="n">ceil_mode</span><span class="p">)</span>

    <span class="k">return</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">out</span><span class="p">)</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">abs</span><span class="p">(</span><span class="n">out</span><span class="p">)))</span><span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">)</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">/</span> <span class="n">norm_type</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">adaptive_max_pool1d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">return_indices</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="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a 1D adaptive max pooling over an input signal composed of</span>
<span class="sd">    several input planes.</span>

<span class="sd">    See :class:`~torch.nn.AdaptiveMaxPool1d` for details and output shape.</span>

<span class="sd">    Args:</span>
<span class="sd">        output_size: the target output size (single integer)</span>
<span class="sd">        return_indices: whether to return pooling indices. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">adaptive_max_pool1d_with_indices</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">adaptive_max_pool1d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_adaptive_max_pool1d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList1</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">return_indices</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">adaptive_max_pool1d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">adaptive_max_pool1d_with_indices</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>


<span class="n">adaptive_max_pool1d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">adaptive_max_pool1d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_adaptive_max_pool1d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;adaptive_max_pool1d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">adaptive_max_pool2d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">return_indices</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="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a 2D adaptive max pooling over an input signal composed of</span>
<span class="sd">    several input planes.</span>

<span class="sd">    See :class:`~torch.nn.AdaptiveMaxPool2d` for details and output shape.</span>

<span class="sd">    Args:</span>
<span class="sd">        output_size: the target output size (single integer or</span>
<span class="sd">            double-integer tuple)</span>
<span class="sd">        return_indices: whether to return pooling indices. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">adaptive_max_pool2d_with_indices</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span>
        <span class="p">)</span>
    <span class="n">output_size</span> <span class="o">=</span> <span class="n">_list_with_default</span><span class="p">(</span><span class="n">output_size</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">adaptive_max_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_adaptive_max_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">return_indices</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">adaptive_max_pool2d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">adaptive_max_pool2d_with_indices</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>


<span class="n">adaptive_max_pool2d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">adaptive_max_pool2d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_adaptive_max_pool2d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;adaptive_max_pool2d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">adaptive_max_pool3d_with_indices</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">return_indices</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="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a 3D adaptive max pooling over an input signal composed of</span>
<span class="sd">    several input planes.</span>

<span class="sd">    See :class:`~torch.nn.AdaptiveMaxPool3d` for details and output shape.</span>

<span class="sd">    Args:</span>
<span class="sd">        output_size: the target output size (single integer or</span>
<span class="sd">            triple-integer tuple)</span>
<span class="sd">        return_indices: whether to return pooling indices. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">adaptive_max_pool3d_with_indices</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span>
        <span class="p">)</span>
    <span class="n">output_size</span> <span class="o">=</span> <span class="n">_list_with_default</span><span class="p">(</span><span class="n">output_size</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">adaptive_max_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_adaptive_max_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">return_indices</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">adaptive_max_pool3d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">return_indices</span><span class="o">=</span><span class="n">return_indices</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">adaptive_max_pool3d_with_indices</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>


<span class="n">adaptive_max_pool3d</span> <span class="o">=</span> <span class="n">boolean_dispatch</span><span class="p">(</span>
    <span class="n">arg_name</span><span class="o">=</span><span class="s2">&quot;return_indices&quot;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
    <span class="n">default</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="n">if_true</span><span class="o">=</span><span class="n">adaptive_max_pool3d_with_indices</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_adaptive_max_pool3d</span><span class="p">,</span>
    <span class="n">module_name</span><span class="o">=</span><span class="vm">__name__</span><span class="p">,</span>
    <span class="n">func_name</span><span class="o">=</span><span class="s2">&quot;adaptive_max_pool3d&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="n">adaptive_avg_pool1d</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">adaptive_avg_pool1d</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">adaptive_avg_pool1d(input, output_size) -&gt; Tensor</span>

<span class="sd">Applies a 1D adaptive average pooling over an input signal composed of</span>
<span class="sd">several input planes.</span>

<span class="sd">See :class:`~torch.nn.AdaptiveAvgPool1d` for details and output shape.</span>

<span class="sd">Args:</span>
<span class="sd">    output_size: the target output size (single integer)</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="adaptive_avg_pool2d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.adaptive_avg_pool2d.html#torch.nn.functional.adaptive_avg_pool2d">[docs]</a><span class="k">def</span> <span class="nf">adaptive_avg_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Applies a 2D adaptive average pooling over an input signal composed of</span>
<span class="sd">    several input planes.</span>

<span class="sd">    See :class:`~torch.nn.AdaptiveAvgPool2d` for details and output shape.</span>

<span class="sd">    Args:</span>
<span class="sd">        output_size: the target output size (single integer or</span>
<span class="sd">            double-integer tuple)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">adaptive_avg_pool2d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="n">_output_size</span> <span class="o">=</span> <span class="n">_list_with_default</span><span class="p">(</span><span class="n">output_size</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">adaptive_avg_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_output_size</span><span class="p">)</span></div>


<div class="viewcode-block" id="adaptive_avg_pool3d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.adaptive_avg_pool3d.html#torch.nn.functional.adaptive_avg_pool3d">[docs]</a><span class="k">def</span> <span class="nf">adaptive_avg_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList3</span><span class="p">[</span><span class="nb">int</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Applies a 3D adaptive average pooling over an input signal composed of</span>
<span class="sd">    several input planes.</span>

<span class="sd">    See :class:`~torch.nn.AdaptiveAvgPool3d` for details and output shape.</span>

<span class="sd">    Args:</span>
<span class="sd">        output_size: the target output size (single integer or</span>
<span class="sd">            triple-integer tuple)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">adaptive_avg_pool3d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="n">_output_size</span> <span class="o">=</span> <span class="n">_list_with_default</span><span class="p">(</span><span class="n">output_size</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">adaptive_avg_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_output_size</span><span class="p">)</span></div>


<span class="c1"># Activation functions</span>
<div class="viewcode-block" id="dropout"><a class="viewcode-back" href="../../../generated/torch.nn.functional.dropout.html#torch.nn.functional.dropout">[docs]</a><span class="k">def</span> <span class="nf">dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    During training, randomly zeroes some of the elements of the input</span>
<span class="sd">    tensor with probability :attr:`p` using samples from a Bernoulli</span>
<span class="sd">    distribution.</span>

<span class="sd">    See :class:`~torch.nn.Dropout` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        p: probability of an element to be zeroed. Default: 0.5</span>
<span class="sd">        training: apply dropout if is ``True``. Default: ``True``</span>
<span class="sd">        inplace: If set to ``True``, will do this operation in-place. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">dropout</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</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">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.0</span> <span class="ow">or</span> <span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;dropout probability has to be between 0 and 1, &quot;</span> <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">p</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">dropout_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">_VF</span><span class="o">.</span><span class="n">dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span></div>


<div class="viewcode-block" id="alpha_dropout"><a class="viewcode-back" href="../../../generated/torch.nn.functional.alpha_dropout.html#torch.nn.functional.alpha_dropout">[docs]</a><span class="k">def</span> <span class="nf">alpha_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies alpha dropout to the input.</span>

<span class="sd">    See :class:`~torch.nn.AlphaDropout` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">alpha_dropout</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</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">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.0</span> <span class="ow">or</span> <span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;dropout probability has to be between 0 and 1, &quot;</span> <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">p</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">alpha_dropout_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">_VF</span><span class="o">.</span><span class="n">alpha_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span></div>


<div class="viewcode-block" id="dropout1d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.dropout1d.html#torch.nn.functional.dropout1d">[docs]</a><span class="k">def</span> <span class="nf">dropout1d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Randomly zero out entire channels (a channel is a 1D feature map,</span>
<span class="sd">    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the</span>
<span class="sd">    batched input is a 1D tensor :math:`\text{input}[i, j]`) of the input tensor).</span>
<span class="sd">    Each channel will be zeroed out independently on every forward call with</span>
<span class="sd">    probability :attr:`p` using samples from a Bernoulli distribution.</span>

<span class="sd">    See :class:`~torch.nn.Dropout1d` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        p: probability of a channel to be zeroed. Default: 0.5</span>
<span class="sd">        training: apply dropout if is ``True``. Default: ``True``</span>
<span class="sd">        inplace: If set to ``True``, will do this operation in-place. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">dropout1d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</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">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.0</span> <span class="ow">or</span> <span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;dropout probability has to be between 0 and 1, &quot;</span> <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">p</span><span class="p">))</span>
    <span class="n">inp_dim</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">inp_dim</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;dropout1d: Expected 2D or 3D input, but received a </span><span class="si">{</span><span class="n">inp_dim</span><span class="si">}</span><span class="s2">D input. &quot;</span>
                           <span class="s2">&quot;Note that dropout1d exists to provide channel-wise dropout on inputs with 1 &quot;</span>
                           <span class="s2">&quot;spatial dimension, a channel dimension, and an optional batch dimension &quot;</span>
                           <span class="s2">&quot;(i.e. 2D or 3D inputs).&quot;</span><span class="p">)</span>

    <span class="n">is_batched</span> <span class="o">=</span> <span class="n">inp_dim</span> <span class="o">==</span> <span class="mi">3</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
        <span class="nb">input</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">unsqueeze_</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="nb">input</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="n">result</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_dropout_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">result</span><span class="o">.</span><span class="n">squeeze_</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">result</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">result</span></div>


<div class="viewcode-block" id="dropout2d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.dropout2d.html#torch.nn.functional.dropout2d">[docs]</a><span class="k">def</span> <span class="nf">dropout2d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Randomly zero out entire channels (a channel is a 2D feature map,</span>
<span class="sd">    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the</span>
<span class="sd">    batched input is a 2D tensor :math:`\text{input}[i, j]`) of the input tensor).</span>
<span class="sd">    Each channel will be zeroed out independently on every forward call with</span>
<span class="sd">    probability :attr:`p` using samples from a Bernoulli distribution.</span>

<span class="sd">    See :class:`~torch.nn.Dropout2d` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        p: probability of a channel to be zeroed. Default: 0.5</span>
<span class="sd">        training: apply dropout if is ``True``. Default: ``True``</span>
<span class="sd">        inplace: If set to ``True``, will do this operation in-place. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">dropout2d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</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">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.0</span> <span class="ow">or</span> <span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;dropout probability has to be between 0 and 1, &quot;</span> <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">p</span><span class="p">))</span>
    <span class="n">inp_dim</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">inp_dim</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">):</span>
        <span class="n">warn_msg</span> <span class="o">=</span> <span class="p">(</span><span class="sa">f</span><span class="s2">&quot;dropout2d: Received a </span><span class="si">{</span><span class="n">inp_dim</span><span class="si">}</span><span class="s2">-D input to dropout2d, which is deprecated &quot;</span>
                    <span class="s2">&quot;and will result in an error in a future release. To retain the behavior &quot;</span>
                    <span class="s2">&quot;and silence this warning, please use dropout instead. Note that dropout2d &quot;</span>
                    <span class="s2">&quot;exists to provide channel-wise dropout on inputs with 2 spatial dimensions, &quot;</span>
                    <span class="s2">&quot;a channel dimension, and an optional batch dimension (i.e. 3D or 4D inputs).&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">warn_msg</span><span class="p">)</span>

    <span class="c1"># TODO: Properly support no-batch-dim inputs. For now, these are NOT supported; passing</span>
    <span class="c1"># a 3D input will perform dropout1d behavior instead. This was done historically and the</span>
    <span class="c1"># behavior is maintained here for now.</span>
    <span class="c1"># See https://github.com/pytorch/pytorch/issues/77081</span>
    <span class="k">if</span> <span class="n">inp_dim</span> <span class="o">==</span> <span class="mi">3</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;dropout2d: Received a 3D input to dropout2d and assuming that channel-wise &quot;</span>
                      <span class="s2">&quot;1D dropout behavior is desired - input is interpreted as shape (N, C, L), where C &quot;</span>
                      <span class="s2">&quot;is the channel dim. This behavior will change in a future release to interpret the &quot;</span>
                      <span class="s2">&quot;input as one without a batch dimension, i.e. shape (C, H, W). To maintain the 1D &quot;</span>
                      <span class="s2">&quot;channel-wise dropout behavior, please switch to using dropout1d instead.&quot;</span><span class="p">)</span>

    <span class="n">result</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_dropout_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">result</span></div>


<div class="viewcode-block" id="dropout3d"><a class="viewcode-back" href="../../../generated/torch.nn.functional.dropout3d.html#torch.nn.functional.dropout3d">[docs]</a><span class="k">def</span> <span class="nf">dropout3d</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Randomly zero out entire channels (a channel is a 3D feature map,</span>
<span class="sd">    e.g., the :math:`j`-th channel of the :math:`i`-th sample in the</span>
<span class="sd">    batched input is a 3D tensor :math:`\text{input}[i, j]`) of the input tensor).</span>
<span class="sd">    Each channel will be zeroed out independently on every forward call with</span>
<span class="sd">    probability :attr:`p` using samples from a Bernoulli distribution.</span>

<span class="sd">    See :class:`~torch.nn.Dropout3d` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        p: probability of a channel to be zeroed. Default: 0.5</span>
<span class="sd">        training: apply dropout if is ``True``. Default: ``True``</span>
<span class="sd">        inplace: If set to ``True``, will do this operation in-place. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">dropout3d</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</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">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.0</span> <span class="ow">or</span> <span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;dropout probability has to be between 0 and 1, &quot;</span> <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">p</span><span class="p">))</span>
    <span class="n">inp_dim</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">inp_dim</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">):</span>
        <span class="n">warn_msg</span> <span class="o">=</span> <span class="p">(</span><span class="sa">f</span><span class="s2">&quot;dropout3d: Received a </span><span class="si">{</span><span class="n">inp_dim</span><span class="si">}</span><span class="s2">-D input to dropout3d, which is deprecated &quot;</span>
                    <span class="s2">&quot;and will result in an error in a future release. To retain the behavior &quot;</span>
                    <span class="s2">&quot;and silence this warning, please use dropout instead. Note that dropout3d &quot;</span>
                    <span class="s2">&quot;exists to provide channel-wise dropout on inputs with 3 spatial dimensions, &quot;</span>
                    <span class="s2">&quot;a channel dimension, and an optional batch dimension (i.e. 4D or 5D inputs).&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">warn_msg</span><span class="p">)</span>

    <span class="n">is_batched</span> <span class="o">=</span> <span class="n">inp_dim</span> <span class="o">==</span> <span class="mi">5</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
        <span class="nb">input</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">unsqueeze_</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="nb">input</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="n">result</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_dropout_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">result</span><span class="o">.</span><span class="n">squeeze_</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">result</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<div class="viewcode-block" id="feature_alpha_dropout"><a class="viewcode-back" href="../../../generated/torch.nn.functional.feature_alpha_dropout.html#torch.nn.functional.feature_alpha_dropout">[docs]</a><span class="k">def</span> <span class="nf">feature_alpha_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.5</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Randomly masks out entire channels (a channel is a feature map,</span>
<span class="sd">    e.g. the :math:`j`-th channel of the :math:`i`-th sample in the batch input</span>
<span class="sd">    is a tensor :math:`\text{input}[i, j]`) of the input tensor). Instead of</span>
<span class="sd">    setting activations to zero, as in regular Dropout, the activations are set</span>
<span class="sd">    to the negative saturation value of the SELU activation function.</span>

<span class="sd">    Each element will be masked independently on every forward call with</span>
<span class="sd">    probability :attr:`p` using samples from a Bernoulli distribution.</span>
<span class="sd">    The elements to be masked are randomized on every forward call, and scaled</span>
<span class="sd">    and shifted to maintain zero mean and unit variance.</span>

<span class="sd">    See :class:`~torch.nn.FeatureAlphaDropout` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        p: dropout probability of a channel to be zeroed. Default: 0.5</span>
<span class="sd">        training: apply dropout if is ``True``. Default: ``True``</span>
<span class="sd">        inplace: If set to ``True``, will do this operation in-place. Default: ``False``</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">feature_alpha_dropout</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</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">inplace</span><span class="o">=</span><span class="n">inplace</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.0</span> <span class="ow">or</span> <span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;dropout probability has to be between 0 and 1, &quot;</span> <span class="s2">&quot;but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">p</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_alpha_dropout_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span> <span class="k">if</span> <span class="n">inplace</span> <span class="k">else</span> <span class="n">_VF</span><span class="o">.</span><span class="n">feature_alpha_dropout</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_threshold</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">threshold</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Thresholds each element of the input Tensor.</span>

<span class="sd">    See :class:`~torch.nn.Threshold` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">_threshold</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">threshold</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">threshold_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">threshold</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">threshold</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">threshold</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span>


<span class="c1"># We define this function as _threshold because it takes an argument</span>
<span class="c1"># named threshold, which clobbers the recursive reference to the</span>
<span class="c1"># function needed for __torch_function__ support</span>
<span class="n">threshold</span> <span class="o">=</span> <span class="n">_threshold</span>

<span class="n">threshold_</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">_VF</span><span class="o">.</span><span class="n">threshold_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">threshold_(input, threshold, value) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~threshold`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="relu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.relu.html#torch.nn.functional.relu">[docs]</a><span class="k">def</span> <span class="nf">relu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;relu(input, inplace=False) -&gt; Tensor</span>

<span class="sd">    Applies the rectified linear unit function element-wise. See</span>
<span class="sd">    :class:`~torch.nn.ReLU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">relu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">relu_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">relu</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">relu_</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">relu_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">relu_(input) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~relu`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="glu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.glu.html#torch.nn.functional.glu">[docs]</a><span class="k">def</span> <span class="nf">glu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    glu(input, dim=-1) -&gt; Tensor</span>

<span class="sd">    The gated linear unit. Computes:</span>

<span class="sd">    .. math ::</span>
<span class="sd">        \text{GLU}(a, b) = a \otimes \sigma(b)</span>

<span class="sd">    where `input` is split in half along `dim` to form `a` and `b`, :math:`\sigma`</span>
<span class="sd">    is the sigmoid function and :math:`\otimes` is the element-wise product between matrices.</span>

<span class="sd">    See `Language Modeling with Gated Convolutional Networks &lt;https://arxiv.org/abs/1612.08083&gt;`_.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input tensor</span>
<span class="sd">        dim (int): dimension on which to split the input. Default: -1</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">glu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;glu does not support scalars because halving size must be even&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">glu</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span></div>


<div class="viewcode-block" id="hardtanh"><a class="viewcode-back" href="../../../generated/torch.nn.functional.hardtanh.html#torch.nn.functional.hardtanh">[docs]</a><span class="k">def</span> <span class="nf">hardtanh</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">min_val</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="o">-</span><span class="mf">1.</span><span class="p">,</span> <span class="n">max_val</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    hardtanh(input, min_val=-1., max_val=1., inplace=False) -&gt; Tensor</span>

<span class="sd">    Applies the HardTanh function element-wise. See :class:`~torch.nn.Hardtanh` for more</span>
<span class="sd">    details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">hardtanh</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">min_val</span><span class="o">=</span><span class="n">min_val</span><span class="p">,</span> <span class="n">max_val</span><span class="o">=</span><span class="n">max_val</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">hardtanh_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">hardtanh</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">min_val</span><span class="p">,</span> <span class="n">max_val</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">hardtanh_</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">hardtanh_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">hardtanh_(input, min_val=-1., max_val=1.) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~hardtanh`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="relu6"><a class="viewcode-back" href="../../../generated/torch.nn.functional.relu6.html#torch.nn.functional.relu6">[docs]</a><span class="k">def</span> <span class="nf">relu6</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;relu6(input, inplace=False) -&gt; Tensor</span>

<span class="sd">    Applies the element-wise function :math:`\text{ReLU6}(x) = \min(\max(0,x), 6)`.</span>

<span class="sd">    See :class:`~torch.nn.ReLU6` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">relu6</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">relu6_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">relu6</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<div class="viewcode-block" id="elu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.elu.html#torch.nn.functional.elu">[docs]</a><span class="k">def</span> <span class="nf">elu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies the Exponential Linear Unit (ELU) function element-wise.</span>

<span class="sd">    See :class:`~torch.nn.ELU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">elu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">elu_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">alpha</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">elu</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">alpha</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">elu_</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">elu_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">elu_(input, alpha=1.) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~elu`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="selu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.selu.html#torch.nn.functional.selu">[docs]</a><span class="k">def</span> <span class="nf">selu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;selu(input, inplace=False) -&gt; Tensor</span>

<span class="sd">    Applies element-wise,</span>
<span class="sd">    :math:`\text{SELU}(x) = scale * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))`,</span>
<span class="sd">    with :math:`\alpha=1.6732632423543772848170429916717` and</span>
<span class="sd">    :math:`scale=1.0507009873554804934193349852946`.</span>

<span class="sd">    See :class:`~torch.nn.SELU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">selu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">selu_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">selu</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">selu_</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">selu_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">selu_(input) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~selu`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="celu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.celu.html#torch.nn.functional.celu">[docs]</a><span class="k">def</span> <span class="nf">celu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;celu(input, alpha=1., inplace=False) -&gt; Tensor</span>

<span class="sd">    Applies element-wise,</span>
<span class="sd">    :math:`\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))`.</span>

<span class="sd">    See :class:`~torch.nn.CELU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">celu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">celu_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">alpha</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">celu</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">alpha</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">celu_</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">celu_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">celu_(input, alpha=1.) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~celu`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="leaky_relu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.leaky_relu.html#torch.nn.functional.leaky_relu">[docs]</a><span class="k">def</span> <span class="nf">leaky_relu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">negative_slope</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.01</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    leaky_relu(input, negative_slope=0.01, inplace=False) -&gt; Tensor</span>

<span class="sd">    Applies element-wise,</span>
<span class="sd">    :math:`\text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)`</span>

<span class="sd">    See :class:`~torch.nn.LeakyReLU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">leaky_relu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">negative_slope</span><span class="o">=</span><span class="n">negative_slope</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">leaky_relu_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">negative_slope</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">leaky_relu</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">negative_slope</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">leaky_relu_</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">leaky_relu_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">leaky_relu_(input, negative_slope=0.01) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~leaky_relu`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="n">prelu</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">prelu</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;prelu(input, weight) -&gt; Tensor</span>

<span class="sd">Applies element-wise the function</span>
<span class="sd">:math:`\text{PReLU}(x) = \max(0,x) + \text{weight} * \min(0,x)` where weight is a</span>
<span class="sd">learnable parameter.</span>

<span class="sd">.. note::</span>
<span class="sd">    `weight` is expected to be a scalar or 1-D tensor. If `weight` is 1-D,</span>
<span class="sd">    its size must match the number of input channels, determined by</span>
<span class="sd">    `input.size(1)` when `input.dim() &gt;= 2`, otherwise 1.</span>
<span class="sd">    In the 1-D case, note that when `input` has dim &gt; 2, `weight` can be expanded</span>
<span class="sd">    to the shape of `input` in a way that is not possible using normal</span>
<span class="sd">    :ref:`broadcasting semantics&lt;broadcasting-semantics&gt;`.</span>

<span class="sd">See :class:`~torch.nn.PReLU` for more details.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>


<div class="viewcode-block" id="rrelu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.rrelu.html#torch.nn.functional.rrelu">[docs]</a><span class="k">def</span> <span class="nf">rrelu</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">lower</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span> <span class="o">/</span> <span class="mi">8</span><span class="p">,</span> <span class="n">upper</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span> <span class="o">/</span> <span class="mi">3</span><span class="p">,</span> <span class="n">training</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">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;rrelu(input, lower=1./8, upper=1./3, training=False, inplace=False) -&gt; Tensor</span>

<span class="sd">    Randomized leaky ReLU.</span>

<span class="sd">    See :class:`~torch.nn.RReLU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">rrelu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">lower</span><span class="o">=</span><span class="n">lower</span><span class="p">,</span> <span class="n">upper</span><span class="o">=</span><span class="n">upper</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">inplace</span><span class="o">=</span><span class="n">inplace</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rrelu_</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">lower</span><span class="p">,</span> <span class="n">upper</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">result</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rrelu</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">lower</span><span class="p">,</span> <span class="n">upper</span><span class="p">,</span> <span class="n">training</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">result</span></div>


<span class="n">rrelu_</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">rrelu_</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">rrelu_(input, lower=1./8, upper=1./3, training=False) -&gt; Tensor</span>

<span class="sd">In-place version of :func:`~rrelu`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">logsigmoid</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">log_sigmoid</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">logsigmoid(input) -&gt; Tensor</span>

<span class="sd">Applies element-wise :math:`\text{LogSigmoid}(x_i) = \log \left(\frac{1}{1 + \exp(-x_i)}\right)`</span>

<span class="sd">See :class:`~torch.nn.LogSigmoid` for more details.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">gelu</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">gelu</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">gelu(input, approximate = &#39;none&#39;) -&gt; Tensor</span>

<span class="sd">When the approximate argument is &#39;none&#39;, it applies element-wise the function</span>
<span class="sd">:math:`\text{GELU}(x) = x * \Phi(x)`</span>

<span class="sd">where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.</span>

<span class="sd">When the approximate argument is &#39;tanh&#39;, Gelu is estimated with</span>

<span class="sd">.. math::</span>
<span class="sd">    \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt(2 / \pi) * (x + 0.044715 * x^3)))</span>

<span class="sd">See `Gaussian Error Linear Units (GELUs) &lt;https://arxiv.org/abs/1606.08415&gt;`_.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>

<span class="n">hardshrink</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">hardshrink</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">hardshrink(input, lambd=0.5) -&gt; Tensor</span>

<span class="sd">Applies the hard shrinkage function element-wise</span>

<span class="sd">See :class:`~torch.nn.Hardshrink` for more details.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>


<div class="viewcode-block" id="tanhshrink"><a class="viewcode-back" href="../../../generated/torch.nn.functional.tanhshrink.html#torch.nn.functional.tanhshrink">[docs]</a><span class="k">def</span> <span class="nf">tanhshrink</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;tanhshrink(input) -&gt; Tensor</span>

<span class="sd">    Applies element-wise, :math:`\text{Tanhshrink}(x) = x - \text{Tanh}(x)`</span>

<span class="sd">    See :class:`~torch.nn.Tanhshrink` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">tanhshrink</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="nb">input</span> <span class="o">-</span> <span class="nb">input</span><span class="o">.</span><span class="n">tanh</span><span class="p">()</span></div>


<div class="viewcode-block" id="softsign"><a class="viewcode-back" href="../../../generated/torch.nn.functional.softsign.html#torch.nn.functional.softsign">[docs]</a><span class="k">def</span> <span class="nf">softsign</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;softsign(input) -&gt; Tensor</span>

<span class="sd">    Applies element-wise, the function :math:`\text{SoftSign}(x) = \frac{x}{1 + |x|}`</span>

<span class="sd">    See :class:`~torch.nn.Softsign` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">softsign</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="nb">input</span> <span class="o">/</span> <span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">abs</span><span class="p">()</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span></div>


<span class="n">softplus</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">softplus</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">softplus(input, beta=1, threshold=20) -&gt; Tensor</span>

<span class="sd">Applies element-wise, the function :math:`\text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))`.</span>

<span class="sd">For numerical stability the implementation reverts to the linear function</span>
<span class="sd">when :math:`input \times \beta &gt; threshold`.</span>

<span class="sd">See :class:`~torch.nn.Softplus` for more details.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="k">def</span> <span class="nf">_get_softmax_dim</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">ndim</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">stacklevel</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">int</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;Implicit dimension choice for </span><span class="si">{}</span><span class="s2"> has been deprecated. &quot;</span>
        <span class="s2">&quot;Change the call to include dim=X as an argument.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">),</span>
        <span class="n">stacklevel</span><span class="o">=</span><span class="n">stacklevel</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="k">if</span> <span class="n">ndim</span> <span class="o">==</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">ndim</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">ndim</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="k">return</span> <span class="n">ret</span>


<div class="viewcode-block" id="softmin"><a class="viewcode-back" href="../../../generated/torch.nn.functional.softmin.html#torch.nn.functional.softmin">[docs]</a><span class="k">def</span> <span class="nf">softmin</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">_stacklevel</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span> <span class="n">dtype</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">DType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a softmin function.</span>

<span class="sd">    Note that :math:`\text{Softmin}(x) = \text{Softmax}(-x)`. See softmax definition for mathematical formula.</span>

<span class="sd">    See :class:`~torch.nn.Softmin` for more details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input</span>
<span class="sd">        dim (int): A dimension along which softmin will be computed (so every slice</span>
<span class="sd">            along dim will sum to 1).</span>
<span class="sd">        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.</span>
<span class="sd">          If specified, the input tensor is casted to :attr:`dtype` before the operation</span>
<span class="sd">          is performed. This is useful for preventing data type overflows. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">softmin</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">,</span> <span class="n">_stacklevel</span><span class="o">=</span><span class="n">_stacklevel</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">dim</span> <span class="o">=</span> <span class="n">_get_softmax_dim</span><span class="p">(</span><span class="s2">&quot;softmin&quot;</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">(),</span> <span class="n">_stacklevel</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dtype</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="nb">input</span><span class="p">)</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="nb">input</span><span class="p">)</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<div class="viewcode-block" id="softmax"><a class="viewcode-back" href="../../../generated/torch.nn.functional.softmax.html#torch.nn.functional.softmax">[docs]</a><span class="k">def</span> <span class="nf">softmax</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">_stacklevel</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span> <span class="n">dtype</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">DType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a softmax function.</span>

<span class="sd">    Softmax is defined as:</span>

<span class="sd">    :math:`\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}`</span>

<span class="sd">    It is applied to all slices along dim, and will re-scale them so that the elements</span>
<span class="sd">    lie in the range `[0, 1]` and sum to 1.</span>

<span class="sd">    See :class:`~torch.nn.Softmax` for more details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input</span>
<span class="sd">        dim (int): A dimension along which softmax will be computed.</span>
<span class="sd">        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.</span>
<span class="sd">          If specified, the input tensor is casted to :attr:`dtype` before the operation</span>
<span class="sd">          is performed. This is useful for preventing data type overflows. Default: None.</span>

<span class="sd">    .. note::</span>
<span class="sd">        This function doesn&#39;t work directly with NLLLoss,</span>
<span class="sd">        which expects the Log to be computed between the Softmax and itself.</span>
<span class="sd">        Use log_softmax instead (it&#39;s faster and has better numerical properties).</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">softmax</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">,</span> <span class="n">_stacklevel</span><span class="o">=</span><span class="n">_stacklevel</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">dim</span> <span class="o">=</span> <span class="n">_get_softmax_dim</span><span class="p">(</span><span class="s2">&quot;softmax&quot;</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">(),</span> <span class="n">_stacklevel</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dtype</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<div class="viewcode-block" id="gumbel_softmax"><a class="viewcode-back" href="../../../generated/torch.nn.functional.gumbel_softmax.html#torch.nn.functional.gumbel_softmax">[docs]</a><span class="k">def</span> <span class="nf">gumbel_softmax</span><span class="p">(</span><span class="n">logits</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">tau</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span> <span class="n">hard</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">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-10</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Samples from the Gumbel-Softmax distribution (`Link 1`_  `Link 2`_) and optionally discretizes.</span>

<span class="sd">    Args:</span>
<span class="sd">      logits: `[..., num_features]` unnormalized log probabilities</span>
<span class="sd">      tau: non-negative scalar temperature</span>
<span class="sd">      hard: if ``True``, the returned samples will be discretized as one-hot vectors,</span>
<span class="sd">            but will be differentiated as if it is the soft sample in autograd</span>
<span class="sd">      dim (int): A dimension along which softmax will be computed. Default: -1.</span>

<span class="sd">    Returns:</span>
<span class="sd">      Sampled tensor of same shape as `logits` from the Gumbel-Softmax distribution.</span>
<span class="sd">      If ``hard=True``, the returned samples will be one-hot, otherwise they will</span>
<span class="sd">      be probability distributions that sum to 1 across `dim`.</span>

<span class="sd">    .. note::</span>
<span class="sd">      This function is here for legacy reasons, may be removed from nn.Functional in the future.</span>

<span class="sd">    .. note::</span>
<span class="sd">      The main trick for `hard` is to do  `y_hard - y_soft.detach() + y_soft`</span>

<span class="sd">      It achieves two things:</span>
<span class="sd">      - makes the output value exactly one-hot</span>
<span class="sd">      (since we add then subtract y_soft value)</span>
<span class="sd">      - makes the gradient equal to y_soft gradient</span>
<span class="sd">      (since we strip all other gradients)</span>

<span class="sd">    Examples::</span>
<span class="sd">        &gt;&gt;&gt; logits = torch.randn(20, 32)</span>
<span class="sd">        &gt;&gt;&gt; # Sample soft categorical using reparametrization trick:</span>
<span class="sd">        &gt;&gt;&gt; F.gumbel_softmax(logits, tau=1, hard=False)</span>
<span class="sd">        &gt;&gt;&gt; # Sample hard categorical using &quot;Straight-through&quot; trick:</span>
<span class="sd">        &gt;&gt;&gt; F.gumbel_softmax(logits, tau=1, hard=True)</span>

<span class="sd">    .. _Link 1:</span>
<span class="sd">        https://arxiv.org/abs/1611.00712</span>
<span class="sd">    .. _Link 2:</span>
<span class="sd">        https://arxiv.org/abs/1611.01144</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="n">logits</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">gumbel_softmax</span><span class="p">,</span> <span class="p">(</span><span class="n">logits</span><span class="p">,),</span> <span class="n">logits</span><span class="p">,</span> <span class="n">tau</span><span class="o">=</span><span class="n">tau</span><span class="p">,</span> <span class="n">hard</span><span class="o">=</span><span class="n">hard</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">eps</span> <span class="o">!=</span> <span class="mf">1e-10</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;`eps` parameter is deprecated and has no effect.&quot;</span><span class="p">)</span>

    <span class="n">gumbels</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">empty_like</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">memory_format</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">legacy_contiguous_format</span><span class="p">)</span><span class="o">.</span><span class="n">exponential_</span><span class="p">()</span><span class="o">.</span><span class="n">log</span><span class="p">()</span>
    <span class="p">)</span>  <span class="c1"># ~Gumbel(0,1)</span>
    <span class="n">gumbels</span> <span class="o">=</span> <span class="p">(</span><span class="n">logits</span> <span class="o">+</span> <span class="n">gumbels</span><span class="p">)</span> <span class="o">/</span> <span class="n">tau</span>  <span class="c1"># ~Gumbel(logits,tau)</span>
    <span class="n">y_soft</span> <span class="o">=</span> <span class="n">gumbels</span><span class="o">.</span><span class="n">softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">hard</span><span class="p">:</span>
        <span class="c1"># Straight through.</span>
        <span class="n">index</span> <span class="o">=</span> <span class="n">y_soft</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">y_hard</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">logits</span><span class="p">,</span> <span class="n">memory_format</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">legacy_contiguous_format</span><span class="p">)</span><span class="o">.</span><span class="n">scatter_</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">index</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">)</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">y_hard</span> <span class="o">-</span> <span class="n">y_soft</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span> <span class="o">+</span> <span class="n">y_soft</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># Reparametrization trick.</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">y_soft</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<div class="viewcode-block" id="log_softmax"><a class="viewcode-back" href="../../../generated/torch.nn.functional.log_softmax.html#torch.nn.functional.log_softmax">[docs]</a><span class="k">def</span> <span class="nf">log_softmax</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">_stacklevel</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">3</span><span class="p">,</span> <span class="n">dtype</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">DType</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies a softmax followed by a logarithm.</span>

<span class="sd">    While mathematically equivalent to log(softmax(x)), doing these two</span>
<span class="sd">    operations separately is slower and numerically unstable. This function</span>
<span class="sd">    uses an alternative formulation to compute the output and gradient correctly.</span>

<span class="sd">    See :class:`~torch.nn.LogSoftmax` for more details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input</span>
<span class="sd">        dim (int): A dimension along which log_softmax will be computed.</span>
<span class="sd">        dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor.</span>
<span class="sd">          If specified, the input tensor is cast to :attr:`dtype` before the operation</span>
<span class="sd">          is performed. This is useful for preventing data type overflows. Default: None.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">log_softmax</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">,</span> <span class="n">_stacklevel</span><span class="o">=</span><span class="n">_stacklevel</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">dim</span> <span class="o">=</span> <span class="n">_get_softmax_dim</span><span class="p">(</span><span class="s2">&quot;log_softmax&quot;</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">(),</span> <span class="n">_stacklevel</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dtype</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">log_softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">log_softmax</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<span class="n">softshrink</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">softshrink</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">softshrink(input, lambd=0.5) -&gt; Tensor</span>

<span class="sd">Applies the soft shrinkage function elementwise</span>

<span class="sd">See :class:`~torch.nn.Softshrink` for more details.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="tanh"><a class="viewcode-back" href="../../../generated/torch.nn.functional.tanh.html#torch.nn.functional.tanh">[docs]</a><span class="k">def</span> <span class="nf">tanh</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;tanh(input) -&gt; Tensor</span>

<span class="sd">    Applies element-wise,</span>
<span class="sd">    :math:`\text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}`</span>

<span class="sd">    See :class:`~torch.nn.Tanh` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="nb">input</span><span class="o">.</span><span class="n">tanh</span><span class="p">()</span></div>


<div class="viewcode-block" id="sigmoid"><a class="viewcode-back" href="../../../generated/torch.nn.functional.sigmoid.html#torch.nn.functional.sigmoid">[docs]</a><span class="k">def</span> <span class="nf">sigmoid</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;sigmoid(input) -&gt; Tensor</span>

<span class="sd">    Applies the element-wise function :math:`\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}`</span>

<span class="sd">    See :class:`~torch.nn.Sigmoid` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="nb">input</span><span class="o">.</span><span class="n">sigmoid</span><span class="p">()</span></div>


<div class="viewcode-block" id="hardsigmoid"><a class="viewcode-back" href="../../../generated/torch.nn.functional.hardsigmoid.html#torch.nn.functional.hardsigmoid">[docs]</a><span class="k">def</span> <span class="nf">hardsigmoid</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies the element-wise function</span>

<span class="sd">    .. math::</span>
<span class="sd">        \text{Hardsigmoid}(x) = \begin{cases}</span>
<span class="sd">            0 &amp; \text{if~} x \le -3, \\</span>
<span class="sd">            1 &amp; \text{if~} x \ge +3, \\</span>
<span class="sd">            x / 6 + 1 / 2 &amp; \text{otherwise}</span>
<span class="sd">        \end{cases}</span>

<span class="sd">    Args:</span>
<span class="sd">        inplace: If set to ``True``, will do this operation in-place. Default: ``False``</span>

<span class="sd">    See :class:`~torch.nn.Hardsigmoid` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">hardsigmoid</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">hardsigmoid_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">hardsigmoid</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span></div>


<span class="n">linear</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">linear</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">linear(input, weight, bias=None) -&gt; Tensor</span>

<span class="sd">Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.</span>

<span class="sd">This operation supports 2-D :attr:`weight` with :ref:`sparse layout&lt;sparse-docs&gt;`</span>

<span class="sd">{sparse_beta_warning}</span>

<span class="sd">This operator supports :ref:`TensorFloat32&lt;tf32_on_ampere&gt;`.</span>

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

<span class="sd">    - Input: :math:`(*, in\_features)` where `*` means any number of</span>
<span class="sd">      additional dimensions, including none</span>
<span class="sd">    - Weight: :math:`(out\_features, in\_features)` or :math:`(in\_features)`</span>
<span class="sd">    - Bias: :math:`(out\_features)` or :math:`()`</span>
<span class="sd">    - Output: :math:`(*, out\_features)` or :math:`(*)`, based on the shape of the weight</span>
<span class="sd">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">sparse_support_notes</span><span class="p">))</span>


<span class="n">bilinear</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">bilinear</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">bilinear(input1, input2, weight, bias=None) -&gt; Tensor</span>

<span class="sd">Applies a bilinear transformation to the incoming data:</span>
<span class="sd">:math:`y = x_1^T A x_2 + b`</span>

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

<span class="sd">    - input1: :math:`(N, *, H_{in1})` where :math:`H_{in1}=\text{in1\_features}`</span>
<span class="sd">      and :math:`*` means any number of additional dimensions.</span>
<span class="sd">      All but the last dimension of the inputs should be the same.</span>
<span class="sd">    - input2: :math:`(N, *, H_{in2})` where :math:`H_{in2}=\text{in2\_features}`</span>
<span class="sd">    - weight: :math:`(\text{out\_features}, \text{in1\_features},</span>
<span class="sd">      \text{in2\_features})`</span>
<span class="sd">    - bias: :math:`(\text{out\_features})`</span>
<span class="sd">    - output: :math:`(N, *, H_{out})` where :math:`H_{out}=\text{out\_features}`</span>
<span class="sd">      and all but the last dimension are the same shape as the input.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>


<div class="viewcode-block" id="silu"><a class="viewcode-back" href="../../../generated/torch.nn.functional.silu.html#torch.nn.functional.silu">[docs]</a><span class="k">def</span> <span class="nf">silu</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies the Sigmoid Linear Unit (SiLU) function, element-wise.</span>
<span class="sd">    The SiLU function is also known as the swish function.</span>

<span class="sd">    .. math::</span>
<span class="sd">        \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}</span>

<span class="sd">    .. note::</span>
<span class="sd">        See `Gaussian Error Linear Units (GELUs) &lt;https://arxiv.org/abs/1606.08415&gt;`_</span>
<span class="sd">        where the SiLU (Sigmoid Linear Unit) was originally coined, and see</span>
<span class="sd">        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation</span>
<span class="sd">        in Reinforcement Learning &lt;https://arxiv.org/abs/1702.03118&gt;`_ and `Swish:</span>
<span class="sd">        a Self-Gated Activation Function &lt;https://arxiv.org/abs/1710.05941v1&gt;`_</span>
<span class="sd">        where the SiLU was experimented with later.</span>

<span class="sd">    See :class:`~torch.nn.SiLU` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">silu</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">silu_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">silu</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span></div>


<div class="viewcode-block" id="mish"><a class="viewcode-back" href="../../../generated/torch.nn.functional.mish.html#torch.nn.functional.mish">[docs]</a><span class="k">def</span> <span class="nf">mish</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies the Mish function, element-wise.</span>
<span class="sd">    Mish: A Self Regularized Non-Monotonic Neural Activation Function.</span>

<span class="sd">    .. math::</span>
<span class="sd">        \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))</span>

<span class="sd">    .. note::</span>
<span class="sd">        See `Mish: A Self Regularized Non-Monotonic Neural Activation Function &lt;https://arxiv.org/abs/1908.08681&gt;`_</span>

<span class="sd">    See :class:`~torch.nn.Mish` for more details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">mish</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">mish_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">mish</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span></div>


<div class="viewcode-block" id="hardswish"><a class="viewcode-back" href="../../../generated/torch.nn.functional.hardswish.html#torch.nn.functional.hardswish">[docs]</a><span class="k">def</span> <span class="nf">hardswish</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">inplace</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies the hardswish function, element-wise, as described in the paper:</span>

<span class="sd">    `Searching for MobileNetV3`_.</span>

<span class="sd">    .. math::</span>
<span class="sd">        \text{Hardswish}(x) = \begin{cases}</span>
<span class="sd">            0 &amp; \text{if~} x \le -3, \\</span>
<span class="sd">            x &amp; \text{if~} x \ge +3, \\</span>
<span class="sd">            x \cdot (x + 3) /6 &amp; \text{otherwise}</span>
<span class="sd">        \end{cases}</span>

<span class="sd">    See :class:`~torch.nn.Hardswish` for more details.</span>

<span class="sd">    .. _`Searching for MobileNetV3`:</span>
<span class="sd">        https://arxiv.org/abs/1905.02244</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">hardswish</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">inplace</span><span class="o">=</span><span class="n">inplace</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">inplace</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">hardswish_</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">hardswish</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_no_grad_embedding_renorm_</span><span class="p">(</span><span class="n">weight</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">max_norm</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">:</span> <span class="nb">float</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">embedding_renorm_</span><span class="p">(</span><span class="n">weight</span><span class="o">.</span><span class="n">detach</span><span class="p">(),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">max_norm</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">)</span>


<div class="viewcode-block" id="embedding"><a class="viewcode-back" href="../../../generated/torch.nn.functional.embedding.html#torch.nn.functional.embedding">[docs]</a><span class="k">def</span> <span class="nf">embedding</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">padding_idx</span><span class="p">:</span> <span class="n">Optional</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="p">,</span>
    <span class="n">max_norm</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">norm_type</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span>
    <span class="n">scale_grad_by_freq</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">sparse</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;A simple lookup table that looks up embeddings in a fixed dictionary and size.</span>

<span class="sd">    This module is often used to retrieve word embeddings using indices.</span>
<span class="sd">    The input to the module is a list of indices, and the embedding matrix,</span>
<span class="sd">    and the output is the corresponding word embeddings.</span>

<span class="sd">    See :class:`torch.nn.Embedding` for more details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (LongTensor): Tensor containing indices into the embedding matrix</span>
<span class="sd">        weight (Tensor): The embedding matrix with number of rows equal to the maximum possible index + 1,</span>
<span class="sd">            and number of columns equal to the embedding size</span>
<span class="sd">        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the gradient;</span>
<span class="sd">                                     therefore, the embedding vector at :attr:`padding_idx` is not updated during training,</span>
<span class="sd">                                     i.e. it remains as a fixed &quot;pad&quot;.</span>
<span class="sd">        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`</span>
<span class="sd">                                    is renormalized to have norm :attr:`max_norm`.</span>
<span class="sd">                                    Note: this will modify :attr:`weight` in-place.</span>
<span class="sd">        norm_type (float, optional): The p of the p-norm to compute for the :attr:`max_norm` option. Default ``2``.</span>
<span class="sd">        scale_grad_by_freq (bool, optional): If given, this will scale gradients by the inverse of frequency of</span>
<span class="sd">                                                the words in the mini-batch. Default ``False``.</span>
<span class="sd">        sparse (bool, optional): If ``True``, gradient w.r.t. :attr:`weight` will be a sparse tensor. See Notes under</span>
<span class="sd">                                 :class:`torch.nn.Embedding` for more details regarding sparse gradients.</span>

<span class="sd">    Shape:</span>
<span class="sd">        - Input: LongTensor of arbitrary shape containing the indices to extract</span>
<span class="sd">        - Weight: Embedding matrix of floating point type with shape `(V, embedding_dim)`,</span>
<span class="sd">          where V = maximum index + 1 and embedding_dim = the embedding size</span>
<span class="sd">        - Output: `(*, embedding_dim)`, where `*` is the input shape</span>

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

<span class="sd">        &gt;&gt;&gt; # a batch of 2 samples of 4 indices each</span>
<span class="sd">        &gt;&gt;&gt; input = torch.tensor([[1, 2, 4, 5], [4, 3, 2, 9]])</span>
<span class="sd">        &gt;&gt;&gt; # an embedding matrix containing 10 tensors of size 3</span>
<span class="sd">        &gt;&gt;&gt; embedding_matrix = torch.rand(10, 3)</span>
<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; F.embedding(input, embedding_matrix)</span>
<span class="sd">        tensor([[[ 0.8490,  0.9625,  0.6753],</span>
<span class="sd">                 [ 0.9666,  0.7761,  0.6108],</span>
<span class="sd">                 [ 0.6246,  0.9751,  0.3618],</span>
<span class="sd">                 [ 0.4161,  0.2419,  0.7383]],</span>

<span class="sd">                [[ 0.6246,  0.9751,  0.3618],</span>
<span class="sd">                 [ 0.0237,  0.7794,  0.0528],</span>
<span class="sd">                 [ 0.9666,  0.7761,  0.6108],</span>
<span class="sd">                 [ 0.3385,  0.8612,  0.1867]]])</span>

<span class="sd">        &gt;&gt;&gt; # example with padding_idx</span>
<span class="sd">        &gt;&gt;&gt; weights = torch.rand(10, 3)</span>
<span class="sd">        &gt;&gt;&gt; weights[0, :].zero_()</span>
<span class="sd">        &gt;&gt;&gt; embedding_matrix = weights</span>
<span class="sd">        &gt;&gt;&gt; input = torch.tensor([[0, 2, 0, 5]])</span>
<span class="sd">        &gt;&gt;&gt; F.embedding(input, embedding_matrix, padding_idx=0)</span>
<span class="sd">        tensor([[[ 0.0000,  0.0000,  0.0000],</span>
<span class="sd">                 [ 0.5609,  0.5384,  0.8720],</span>
<span class="sd">                 [ 0.0000,  0.0000,  0.0000],</span>
<span class="sd">                 [ 0.6262,  0.2438,  0.7471]]])</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">embedding</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">weight</span><span class="p">,</span>
            <span class="n">padding_idx</span><span class="o">=</span><span class="n">padding_idx</span><span class="p">,</span>
            <span class="n">max_norm</span><span class="o">=</span><span class="n">max_norm</span><span class="p">,</span>
            <span class="n">norm_type</span><span class="o">=</span><span class="n">norm_type</span><span class="p">,</span>
            <span class="n">scale_grad_by_freq</span><span class="o">=</span><span class="n">scale_grad_by_freq</span><span class="p">,</span>
            <span class="n">sparse</span><span class="o">=</span><span class="n">sparse</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">padding_idx</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">padding_idx</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">padding_idx</span> <span class="o">&lt;</span> <span class="n">weight</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="s2">&quot;Padding_idx must be within num_embeddings&quot;</span>
        <span class="k">elif</span> <span class="n">padding_idx</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">padding_idx</span> <span class="o">&gt;=</span> <span class="o">-</span><span class="n">weight</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">),</span> <span class="s2">&quot;Padding_idx must be within num_embeddings&quot;</span>
            <span class="n">padding_idx</span> <span class="o">=</span> <span class="n">weight</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">+</span> <span class="n">padding_idx</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">padding_idx</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
    <span class="k">if</span> <span class="n">max_norm</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># Note [embedding_renorm contiguous]</span>
        <span class="c1"># `embedding_renorm_` will call .contiguous() on input anyways, so we</span>
        <span class="c1"># call it here and take advantage of the improved locality in the</span>
        <span class="c1"># `embedding` call below too.</span>
        <span class="nb">input</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span>
        <span class="c1"># Note [embedding_renorm set_grad_enabled]</span>
        <span class="c1"># XXX: equivalent to</span>
        <span class="c1"># with torch.no_grad():</span>
        <span class="c1">#   torch.embedding_renorm_</span>
        <span class="c1"># remove once script supports set_grad_enabled</span>
        <span class="n">_no_grad_embedding_renorm_</span><span class="p">(</span><span class="n">weight</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">max_norm</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">embedding</span><span class="p">(</span><span class="n">weight</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">padding_idx</span><span class="p">,</span> <span class="n">scale_grad_by_freq</span><span class="p">,</span> <span class="n">sparse</span><span class="p">)</span></div>


<div class="viewcode-block" id="embedding_bag"><a class="viewcode-back" href="../../../generated/torch.nn.functional.embedding_bag.html#torch.nn.functional.embedding_bag">[docs]</a><span class="k">def</span> <span class="nf">embedding_bag</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">offsets</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">max_norm</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">norm_type</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">2</span><span class="p">,</span>
    <span class="n">scale_grad_by_freq</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">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
    <span class="n">sparse</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">per_sample_weights</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">include_last_offset</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">padding_idx</span><span class="p">:</span> <span class="n">Optional</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="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Computes sums, means or maxes of `bags` of embeddings, without instantiating the</span>
<span class="sd">    intermediate embeddings.</span>

<span class="sd">    See :class:`torch.nn.EmbeddingBag` for more details.</span>

<span class="sd">    Note:</span>
<span class="sd">        {backward_reproducibility_note}</span>

<span class="sd">    Args:</span>
<span class="sd">        input (LongTensor): Tensor containing bags of indices into the embedding matrix</span>
<span class="sd">        weight (Tensor): The embedding matrix with number of rows equal to the maximum possible index + 1,</span>
<span class="sd">            and number of columns equal to the embedding size</span>
<span class="sd">        offsets (LongTensor, optional): Only used when :attr:`input` is 1D. :attr:`offsets` determines</span>
<span class="sd">                             the starting index position of each bag (sequence) in :attr:`input`.</span>
<span class="sd">        max_norm (float, optional): If given, each embedding vector with norm larger than :attr:`max_norm`</span>
<span class="sd">                                    is renormalized to have norm :attr:`max_norm`.</span>
<span class="sd">                                    Note: this will modify :attr:`weight` in-place.</span>
<span class="sd">        norm_type (float, optional): The ``p`` in the ``p``-norm to compute for the :attr:`max_norm` option.</span>
<span class="sd">                                     Default ``2``.</span>
<span class="sd">        scale_grad_by_freq (bool, optional): if given, this will scale gradients by the inverse of frequency of</span>
<span class="sd">                                                the words in the mini-batch. Default ``False``.</span>
<span class="sd">                                                Note: this option is not supported when ``mode=&quot;max&quot;``.</span>
<span class="sd">        mode (str, optional): ``&quot;sum&quot;``, ``&quot;mean&quot;`` or ``&quot;max&quot;``. Specifies the way to reduce the bag.</span>
<span class="sd">                                 Default: ``&quot;mean&quot;``</span>
<span class="sd">        sparse (bool, optional): if ``True``, gradient w.r.t. :attr:`weight` will be a sparse tensor. See Notes under</span>
<span class="sd">                                 :class:`torch.nn.Embedding` for more details regarding sparse gradients.</span>
<span class="sd">                                 Note: this option is not supported when ``mode=&quot;max&quot;``.</span>
<span class="sd">        per_sample_weights (Tensor, optional): a tensor of float / double weights, or None</span>
<span class="sd">            to indicate all weights should be taken to be 1. If specified, :attr:`per_sample_weights`</span>
<span class="sd">            must have exactly the same shape as input and is treated as having the same</span>
<span class="sd">            :attr:`offsets`, if those are not None.</span>

<span class="sd">        include_last_offset (bool, optional): if ``True``, the size of offsets is equal to the number of bags + 1.</span>
<span class="sd">            The last element is the size of the input, or the ending index position of the last bag (sequence).</span>

<span class="sd">        padding_idx (int, optional): If specified, the entries at :attr:`padding_idx` do not contribute to the</span>
<span class="sd">                                     gradient; therefore, the embedding vector at :attr:`padding_idx` is not updated</span>
<span class="sd">                                     during training, i.e. it remains as a fixed &quot;pad&quot;. Note that the embedding</span>
<span class="sd">                                     vector at :attr:`padding_idx` is excluded from the reduction.</span>

<span class="sd">    Shape:</span>
<span class="sd">        - :attr:`input` (LongTensor) and :attr:`offsets` (LongTensor, optional)</span>

<span class="sd">          - If :attr:`input` is 2D of shape `(B, N)`, it will be treated as ``B`` bags (sequences)</span>
<span class="sd">            each of fixed length ``N``, and this will return ``B`` values aggregated in a way</span>
<span class="sd">            depending on the :attr:`mode`. :attr:`offsets` is ignored and required to be ``None`` in this case.</span>

<span class="sd">          - If :attr:`input` is 1D of shape `(N)`, it will be treated as a concatenation of</span>
<span class="sd">            multiple bags (sequences). :attr:`offsets` is required to be a 1D tensor containing</span>
<span class="sd">            the starting index positions of each bag in :attr:`input`. Therefore, for :attr:`offsets`</span>
<span class="sd">            of shape `(B)`, :attr:`input` will be viewed as having ``B`` bags.</span>
<span class="sd">            Empty bags (i.e., having 0-length) will have returned vectors filled by zeros.</span>

<span class="sd">        - :attr:`weight` (Tensor): the learnable weights of the module of shape `(num_embeddings, embedding_dim)`</span>

<span class="sd">        - :attr:`per_sample_weights` (Tensor, optional). Has the same shape as :attr:`input`.</span>

<span class="sd">        - :attr:`output`: aggregated embedding values of shape `(B, embedding_dim)`</span>

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

<span class="sd">        &gt;&gt;&gt; # an Embedding module containing 10 tensors of size 3</span>
<span class="sd">        &gt;&gt;&gt; embedding_matrix = torch.rand(10, 3)</span>
<span class="sd">        &gt;&gt;&gt; # a batch of 2 samples of 4 indices each</span>
<span class="sd">        &gt;&gt;&gt; input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9])</span>
<span class="sd">        &gt;&gt;&gt; offsets = torch.tensor([0, 4])</span>
<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; F.embedding_bag(input, embedding_matrix, offsets)</span>
<span class="sd">        tensor([[ 0.3397,  0.3552,  0.5545],</span>
<span class="sd">                [ 0.5893,  0.4386,  0.5882]])</span>

<span class="sd">        &gt;&gt;&gt; # example with padding_idx</span>
<span class="sd">        &gt;&gt;&gt; embedding_matrix = torch.rand(10, 3)</span>
<span class="sd">        &gt;&gt;&gt; input = torch.tensor([2, 2, 2, 2, 4, 3, 2, 9])</span>
<span class="sd">        &gt;&gt;&gt; offsets = torch.tensor([0, 4])</span>
<span class="sd">        &gt;&gt;&gt; F.embedding_bag(input, embedding_matrix, offsets, padding_idx=2, mode=&#39;sum&#39;)</span>
<span class="sd">        tensor([[ 0.0000,  0.0000,  0.0000],</span>
<span class="sd">                [-0.7082,  3.2145, -2.6251]])</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">offsets</span><span class="p">,</span> <span class="n">per_sample_weights</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">embedding_bag</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">offsets</span><span class="p">,</span> <span class="n">per_sample_weights</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">weight</span><span class="p">,</span>
            <span class="n">offsets</span><span class="o">=</span><span class="n">offsets</span><span class="p">,</span>
            <span class="n">max_norm</span><span class="o">=</span><span class="n">max_norm</span><span class="p">,</span>
            <span class="n">norm_type</span><span class="o">=</span><span class="n">norm_type</span><span class="p">,</span>
            <span class="n">scale_grad_by_freq</span><span class="o">=</span><span class="n">scale_grad_by_freq</span><span class="p">,</span>
            <span class="n">mode</span><span class="o">=</span><span class="n">mode</span><span class="p">,</span>
            <span class="n">sparse</span><span class="o">=</span><span class="n">sparse</span><span class="p">,</span>
            <span class="n">per_sample_weights</span><span class="o">=</span><span class="n">per_sample_weights</span><span class="p">,</span>
            <span class="n">include_last_offset</span><span class="o">=</span><span class="n">include_last_offset</span><span class="p">,</span>
            <span class="n">padding_idx</span><span class="o">=</span><span class="n">padding_idx</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="c1"># Check for backward compatibility.</span>
    <span class="c1"># Used to be embedding_bag(weight, input, ...)</span>
    <span class="c1"># Now is     embedding_bag(input, weight, ...)</span>
    <span class="k">if</span> <span class="n">weight</span><span class="o">.</span><span class="n">dtype</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">long</span> <span class="ow">and</span> <span class="nb">input</span><span class="o">.</span><span class="n">is_floating_point</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;Argument order of nn.functional.embedding_bag was changed. &quot;</span>
            <span class="s2">&quot;Usage `embedding_bag(weight, input, ...)` is deprecated, &quot;</span>
            <span class="s2">&quot;and should now be `embedding_bag(input, weight, ...)`.&quot;</span>
        <span class="p">)</span>
        <span class="n">weight</span><span class="p">,</span> <span class="nb">input</span> <span class="o">=</span> <span class="nb">input</span><span class="p">,</span> <span class="n">weight</span>

    <span class="k">if</span> <span class="n">per_sample_weights</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">!=</span> <span class="n">per_sample_weights</span><span class="o">.</span><span class="n">size</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;embedding_bag: If per_sample_weights (</span><span class="si">{}</span><span class="s2">) is not None, &quot;</span>
            <span class="s2">&quot;then it must have the same shape as the input (</span><span class="si">{}</span><span class="s2">)&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">per_sample_weights</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">offsets</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">type_str</span> <span class="o">=</span> <span class="s2">&quot;&lt;unknown&gt;&quot;</span>
            <span class="c1"># TODO: Remove this once script supports type() calls</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">():</span>
                <span class="n">type_str</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">offsets</span><span class="p">))</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;if input is 2D, then offsets has to be None&quot;</span>
                <span class="s2">&quot;, as input is treated is a mini-batch of&quot;</span>
                <span class="s2">&quot; fixed length sequences. However, found &quot;</span>
                <span class="s2">&quot;offsets of type </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">type_str</span><span class="p">)</span>
            <span class="p">)</span>
        <span class="n">offsets</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">input</span><span class="o">.</span><span class="n">numel</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="nb">input</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="nb">input</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>

        <span class="nb">input</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">reshape</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">per_sample_weights</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">per_sample_weights</span> <span class="o">=</span> <span class="n">per_sample_weights</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
    <span class="k">elif</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</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">offsets</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;offsets has to be a 1D Tensor but got None&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">offsets</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;offsets has to be a 1D Tensor&quot;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;input has to be 1D or 2D Tensor,&quot;</span> <span class="s2">&quot; but got Tensor of dimension </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()))</span>
    <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;sum&quot;</span><span class="p">:</span>
        <span class="n">mode_enum</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">elif</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;mean&quot;</span><span class="p">:</span>
        <span class="n">mode_enum</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="k">elif</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;max&quot;</span><span class="p">:</span>
        <span class="n">mode_enum</span> <span class="o">=</span> <span class="mi">2</span>

        <span class="k">if</span> <span class="n">scale_grad_by_freq</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;max mode does not support scaling the gradient by the frequency&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">sparse</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;max mode does not support sparse weights&quot;</span><span class="p">)</span>

    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;mode has to be one of sum, mean or max&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">max_norm</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># XXX: equivalent to</span>
        <span class="c1"># with torch.no_grad():</span>
        <span class="c1">#   torch.nembedding_renorm_</span>
        <span class="c1"># remove once script supports set_grad_enabled</span>
        <span class="n">_no_grad_embedding_renorm_</span><span class="p">(</span><span class="n">weight</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">max_norm</span><span class="p">,</span> <span class="n">norm_type</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">per_sample_weights</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">!=</span> <span class="s2">&quot;sum&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span>
            <span class="s2">&quot;embedding_bag: per_sample_weights was not None. &quot;</span>
            <span class="s2">&quot;per_sample_weights is only supported for mode=&#39;sum&#39; &quot;</span>
            <span class="s2">&quot;(got mode=&#39;</span><span class="si">{}</span><span class="s2">&#39;). Please open a feature request on GitHub.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span>
        <span class="p">)</span>

    <span class="n">ret</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</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">embedding_bag</span><span class="p">(</span>
        <span class="n">weight</span><span class="p">,</span> <span class="nb">input</span><span class="p">,</span> <span class="n">offsets</span><span class="p">,</span> <span class="n">scale_grad_by_freq</span><span class="p">,</span> <span class="n">mode_enum</span><span class="p">,</span> <span class="n">sparse</span><span class="p">,</span> <span class="n">per_sample_weights</span><span class="p">,</span> <span class="n">include_last_offset</span><span class="p">,</span> <span class="n">padding_idx</span>
    <span class="p">)</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<span class="k">if</span> <span class="n">embedding_bag</span><span class="o">.</span><span class="vm">__doc__</span><span class="p">:</span>
    <span class="n">embedding_bag</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">embedding_bag</span><span class="o">.</span><span class="vm">__doc__</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_verify_batch_size</span><span class="p">(</span><span class="n">size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
    <span class="c1"># XXX: JIT script does not support the reduce from functools, and mul op is a</span>
    <span class="c1"># builtin, which cannot be used as a value to a func yet, so rewrite this size</span>
    <span class="c1"># check to a simple equivalent for loop</span>
    <span class="c1">#</span>
    <span class="c1"># TODO: make use of reduce like below when JIT is ready with the missing features:</span>
    <span class="c1"># from operator import mul</span>
    <span class="c1"># from functools import reduce</span>
    <span class="c1">#</span>
    <span class="c1">#   if reduce(mul, size[2:], size[0]) == 1</span>
    <span class="n">size_prods</span> <span class="o">=</span> <span class="n">size</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">size</span><span class="p">)</span> <span class="o">-</span> <span class="mi">2</span><span class="p">):</span>
        <span class="n">size_prods</span> <span class="o">*=</span> <span class="n">size</span><span class="p">[</span><span class="n">i</span> <span class="o">+</span> <span class="mi">2</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">size_prods</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Expected more than 1 value per channel when training, got input size </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">size</span><span class="p">))</span>


<div class="viewcode-block" id="batch_norm"><a class="viewcode-back" href="../../../generated/torch.nn.functional.batch_norm.html#torch.nn.functional.batch_norm">[docs]</a><span class="k">def</span> <span class="nf">batch_norm</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">running_mean</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">running_var</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">bias</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">training</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">momentum</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
    <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies Batch Normalization for each channel across a batch of data.</span>

<span class="sd">    See :class:`~torch.nn.BatchNorm1d`, :class:`~torch.nn.BatchNorm2d`,</span>
<span class="sd">    :class:`~torch.nn.BatchNorm3d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">running_mean</span><span class="p">,</span> <span class="n">running_var</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">batch_norm</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">running_mean</span><span class="p">,</span> <span class="n">running_var</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">running_mean</span><span class="p">,</span>
            <span class="n">running_var</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">bias</span><span class="o">=</span><span class="n">bias</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">momentum</span><span class="o">=</span><span class="n">momentum</span><span class="p">,</span>
            <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">training</span><span class="p">:</span>
        <span class="n">_verify_batch_size</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">batch_norm</span><span class="p">(</span>
        <span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,</span> <span class="n">running_mean</span><span class="p">,</span> <span class="n">running_var</span><span class="p">,</span> <span class="n">training</span><span class="p">,</span> <span class="n">momentum</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">backends</span><span class="o">.</span><span class="n">cudnn</span><span class="o">.</span><span class="n">enabled</span>
    <span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_verify_spatial_size</span><span class="p">(</span><span class="n">size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
    <span class="c1"># Verify that there is &gt; 1 spatial element for instance norm calculation.</span>
    <span class="n">size_prods</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">size</span><span class="p">)):</span>
        <span class="n">size_prods</span> <span class="o">*=</span> <span class="n">size</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">size_prods</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Expected more than 1 spatial element when training, got input size </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">size</span><span class="p">))</span>


<div class="viewcode-block" id="instance_norm"><a class="viewcode-back" href="../../../generated/torch.nn.functional.instance_norm.html#torch.nn.functional.instance_norm">[docs]</a><span class="k">def</span> <span class="nf">instance_norm</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">running_mean</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">running_var</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">bias</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">use_input_stats</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">momentum</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.1</span><span class="p">,</span>
    <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies Instance Normalization for each channel in each data sample in a</span>
<span class="sd">    batch.</span>

<span class="sd">    See :class:`~torch.nn.InstanceNorm1d`, :class:`~torch.nn.InstanceNorm2d`,</span>
<span class="sd">    :class:`~torch.nn.InstanceNorm3d` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">running_mean</span><span class="p">,</span> <span class="n">running_var</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">instance_norm</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">running_mean</span><span class="p">,</span> <span class="n">running_var</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">running_mean</span><span class="o">=</span><span class="n">running_mean</span><span class="p">,</span>
            <span class="n">running_var</span><span class="o">=</span><span class="n">running_var</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">bias</span><span class="o">=</span><span class="n">bias</span><span class="p">,</span>
            <span class="n">use_input_stats</span><span class="o">=</span><span class="n">use_input_stats</span><span class="p">,</span>
            <span class="n">momentum</span><span class="o">=</span><span class="n">momentum</span><span class="p">,</span>
            <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">use_input_stats</span><span class="p">:</span>
        <span class="n">_verify_spatial_size</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">instance_norm</span><span class="p">(</span>
        <span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,</span> <span class="n">running_mean</span><span class="p">,</span> <span class="n">running_var</span><span class="p">,</span> <span class="n">use_input_stats</span><span class="p">,</span> <span class="n">momentum</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">backends</span><span class="o">.</span><span class="n">cudnn</span><span class="o">.</span><span class="n">enabled</span>
    <span class="p">)</span></div>


<div class="viewcode-block" id="layer_norm"><a class="viewcode-back" href="../../../generated/torch.nn.functional.layer_norm.html#torch.nn.functional.layer_norm">[docs]</a><span class="k">def</span> <span class="nf">layer_norm</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">normalized_shape</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">bias</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies Layer Normalization for last certain number of dimensions.</span>

<span class="sd">    See :class:`~torch.nn.LayerNorm` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">layer_norm</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">normalized_shape</span><span class="p">,</span> <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="o">=</span><span class="n">bias</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">layer_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">normalized_shape</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">backends</span><span class="o">.</span><span class="n">cudnn</span><span class="o">.</span><span class="n">enabled</span><span class="p">)</span></div>


<div class="viewcode-block" id="group_norm"><a class="viewcode-back" href="../../../generated/torch.nn.functional.group_norm.html#torch.nn.functional.group_norm">[docs]</a><span class="k">def</span> <span class="nf">group_norm</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">num_groups</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">bias</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies Group Normalization for last certain number of dimensions.</span>

<span class="sd">    See :class:`~torch.nn.GroupNorm` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">group_norm</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">num_groups</span><span class="p">,</span> <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="o">=</span><span class="n">bias</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Expected at least 2 dimensions for input tensor but received </span><span class="si">{</span><span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="n">_verify_batch_size</span><span class="p">([</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">*</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span> <span class="o">//</span> <span class="n">num_groups</span><span class="p">,</span> <span class="n">num_groups</span><span class="p">]</span> <span class="o">+</span> <span class="nb">list</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()[</span><span class="mi">2</span><span class="p">:]))</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">group_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">num_groups</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">bias</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">backends</span><span class="o">.</span><span class="n">cudnn</span><span class="o">.</span><span class="n">enabled</span><span class="p">)</span></div>


<div class="viewcode-block" id="local_response_norm"><a class="viewcode-back" href="../../../generated/torch.nn.functional.local_response_norm.html#torch.nn.functional.local_response_norm">[docs]</a><span class="k">def</span> <span class="nf">local_response_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-4</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.75</span><span class="p">,</span> <span class="n">k</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Applies local response normalization over an input signal composed of</span>
<span class="sd">    several input planes, where channels occupy the second dimension.</span>
<span class="sd">    Applies normalization across channels.</span>

<span class="sd">    See :class:`~torch.nn.LocalResponseNorm` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">local_response_norm</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span> <span class="n">beta</span><span class="o">=</span><span class="n">beta</span><span class="p">,</span> <span class="n">k</span><span class="o">=</span><span class="n">k</span><span class="p">)</span>
    <span class="n">dim</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">dim</span> <span class="o">&lt;</span> <span class="mi">3</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;Expected 3D or higher dimensionality </span><span class="se">\</span>
<span class="s2">                         input (got </span><span class="si">{}</span><span class="s2"> dimensions)&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span>
                <span class="n">dim</span>
            <span class="p">)</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">numel</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="nb">input</span>

    <span class="n">div</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">dim</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">div</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">pad</span><span class="p">(</span><span class="n">div</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">size</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="p">(</span><span class="n">size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="o">//</span> <span class="mi">2</span><span class="p">))</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">avg_pool2d</span><span class="p">(</span><span class="n">div</span><span class="p">,</span> <span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">stride</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">sizes</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">div</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">sizes</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="mi">1</span><span class="p">,</span> <span class="n">sizes</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">sizes</span><span class="p">[</span><span class="mi">2</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">pad</span><span class="p">(</span><span class="n">div</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">size</span> <span class="o">//</span> <span class="mi">2</span><span class="p">,</span> <span class="p">(</span><span class="n">size</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span> <span class="o">//</span> <span class="mi">2</span><span class="p">))</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">avg_pool3d</span><span class="p">(</span><span class="n">div</span><span class="p">,</span> <span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="n">stride</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">div</span> <span class="o">=</span> <span class="n">div</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">sizes</span><span class="p">)</span>
    <span class="n">div</span> <span class="o">=</span> <span class="n">div</span><span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="n">alpha</span><span class="p">)</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">k</span><span class="p">)</span><span class="o">.</span><span class="n">pow</span><span class="p">(</span><span class="n">beta</span><span class="p">)</span>
    <span class="k">return</span> <span class="nb">input</span> <span class="o">/</span> <span class="n">div</span></div>


<span class="c1"># loss</span>


<div class="viewcode-block" id="ctc_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.ctc_loss.html#torch.nn.functional.ctc_loss">[docs]</a><span class="k">def</span> <span class="nf">ctc_loss</span><span class="p">(</span>
    <span class="n">log_probs</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">targets</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">input_lengths</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target_lengths</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">blank</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
    <span class="n">zero_infinity</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;The Connectionist Temporal Classification loss.</span>

<span class="sd">    See :class:`~torch.nn.CTCLoss` for details.</span>

<span class="sd">    Note:</span>
<span class="sd">        {cudnn_reproducibility_note}</span>

<span class="sd">    Note:</span>
<span class="sd">        {backward_reproducibility_note}</span>

<span class="sd">    Args:</span>
<span class="sd">        log_probs: :math:`(T, N, C)` or :math:`(T, C)` where `C = number of characters in alphabet including blank`,</span>
<span class="sd">            `T = input length`, and `N = batch size`.</span>
<span class="sd">            The logarithmized probabilities of the outputs</span>
<span class="sd">            (e.g. obtained with :func:`torch.nn.functional.log_softmax`).</span>
<span class="sd">        targets: :math:`(N, S)` or `(sum(target_lengths))`.</span>
<span class="sd">            Targets cannot be blank. In the second form, the targets are assumed to be concatenated.</span>
<span class="sd">        input_lengths: :math:`(N)` or :math:`()`.</span>
<span class="sd">            Lengths of the inputs (must each be :math:`\leq T`)</span>
<span class="sd">        target_lengths: :math:`(N)` or :math:`()`.</span>
<span class="sd">            Lengths of the targets</span>
<span class="sd">        blank (int, optional):</span>
<span class="sd">            Blank label. Default :math:`0`.</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the output losses will be divided by the target lengths and</span>
<span class="sd">            then the mean over the batch is taken, ``&#39;sum&#39;``: the output will be</span>
<span class="sd">            summed. Default: ``&#39;mean&#39;``</span>
<span class="sd">        zero_infinity (bool, optional):</span>
<span class="sd">            Whether to zero infinite losses and the associated gradients.</span>
<span class="sd">            Default: ``False``</span>
<span class="sd">            Infinite losses mainly occur when the inputs are too short</span>
<span class="sd">            to be aligned to the targets.</span>

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

<span class="sd">        &gt;&gt;&gt; log_probs = torch.randn(50, 16, 20).log_softmax(2).detach().requires_grad_()</span>
<span class="sd">        &gt;&gt;&gt; targets = torch.randint(1, 20, (16, 30), dtype=torch.long)</span>
<span class="sd">        &gt;&gt;&gt; input_lengths = torch.full((16,), 50, dtype=torch.long)</span>
<span class="sd">        &gt;&gt;&gt; target_lengths = torch.randint(10, 30, (16,), dtype=torch.long)</span>
<span class="sd">        &gt;&gt;&gt; loss = F.ctc_loss(log_probs, targets, input_lengths, target_lengths)</span>
<span class="sd">        &gt;&gt;&gt; loss.backward()</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="n">log_probs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">input_lengths</span><span class="p">,</span> <span class="n">target_lengths</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">ctc_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="n">log_probs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">input_lengths</span><span class="p">,</span> <span class="n">target_lengths</span><span class="p">),</span>
            <span class="n">log_probs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">input_lengths</span><span class="p">,</span> <span class="n">target_lengths</span><span class="p">,</span>
            <span class="n">blank</span><span class="o">=</span><span class="n">blank</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span> <span class="n">zero_infinity</span><span class="o">=</span><span class="n">zero_infinity</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">ctc_loss</span><span class="p">(</span>
        <span class="n">log_probs</span><span class="p">,</span> <span class="n">targets</span><span class="p">,</span> <span class="n">input_lengths</span><span class="p">,</span> <span class="n">target_lengths</span><span class="p">,</span> <span class="n">blank</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">),</span> <span class="n">zero_infinity</span>
    <span class="p">)</span></div>


<span class="k">if</span> <span class="n">ctc_loss</span><span class="o">.</span><span class="vm">__doc__</span><span class="p">:</span>
    <span class="n">ctc_loss</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">ctc_loss</span><span class="o">.</span><span class="vm">__doc__</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>


<div class="viewcode-block" id="nll_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.nll_loss.html#torch.nn.functional.nll_loss">[docs]</a><span class="k">def</span> <span class="nf">nll_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">ignore_index</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;The negative log likelihood loss.</span>

<span class="sd">    See :class:`~torch.nn.NLLLoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: :math:`(N, C)` where `C = number of classes` or :math:`(N, C, H, W)`</span>
<span class="sd">            in case of 2D Loss, or :math:`(N, C, d_1, d_2, ..., d_K)` where :math:`K \geq 1`</span>
<span class="sd">            in the case of K-dimensional loss. `input` is expected to be log-probabilities.</span>
<span class="sd">        target: :math:`(N)` where each value is :math:`0 \leq \text{targets}[i] \leq C-1`,</span>
<span class="sd">            or :math:`(N, d_1, d_2, ..., d_K)` where :math:`K \geq 1` for</span>
<span class="sd">            K-dimensional loss.</span>
<span class="sd">        weight (Tensor, optional): a manual rescaling weight given to each</span>
<span class="sd">            class. If given, has to be a Tensor of size `C`</span>
<span class="sd">        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,</span>
<span class="sd">            the losses are averaged over each loss element in the batch. Note that for</span>
<span class="sd">            some losses, there multiple elements per sample. If the field :attr:`size_average`</span>
<span class="sd">            is set to ``False``, the losses are instead summed for each minibatch. Ignored</span>
<span class="sd">            when reduce is ``False``. Default: ``True``</span>
<span class="sd">        ignore_index (int, optional): Specifies a target value that is ignored</span>
<span class="sd">            and does not contribute to the input gradient. When :attr:`size_average` is</span>
<span class="sd">            ``True``, the loss is averaged over non-ignored targets. Default: -100</span>
<span class="sd">        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the</span>
<span class="sd">            losses are averaged or summed over observations for each minibatch depending</span>
<span class="sd">            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per</span>
<span class="sd">            batch element instead and ignores :attr:`size_average`. Default: ``True``</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the sum of the output will be divided by the number of</span>
<span class="sd">            elements in the output, ``&#39;sum&#39;``: the output will be summed. Note: :attr:`size_average`</span>
<span class="sd">            and :attr:`reduce` are in the process of being deprecated, and in the meantime,</span>
<span class="sd">            specifying either of those two args will override :attr:`reduction`. Default: ``&#39;mean&#39;``</span>

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

<span class="sd">        &gt;&gt;&gt; # input is of size N x C = 3 x 5</span>
<span class="sd">        &gt;&gt;&gt; input = torch.randn(3, 5, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; # each element in target has to have 0 &lt;= value &lt; C</span>
<span class="sd">        &gt;&gt;&gt; target = torch.tensor([1, 0, 4])</span>
<span class="sd">        &gt;&gt;&gt; output = F.nll_loss(F.log_softmax(input, dim=1), target)</span>
<span class="sd">        &gt;&gt;&gt; output.backward()</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">nll_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">nll_loss_nd</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">),</span> <span class="n">ignore_index</span><span class="p">)</span></div>


<div class="viewcode-block" id="poisson_nll_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.poisson_nll_loss.html#torch.nn.functional.poisson_nll_loss">[docs]</a><span class="k">def</span> <span class="nf">poisson_nll_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">log_input</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">full</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">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-8</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Poisson negative log likelihood loss.</span>

<span class="sd">    See :class:`~torch.nn.PoissonNLLLoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: expectation of underlying Poisson distribution.</span>
<span class="sd">        target: random sample :math:`target \sim \text{Poisson}(input)`.</span>
<span class="sd">        log_input: if ``True`` the loss is computed as</span>
<span class="sd">            :math:`\exp(\text{input}) - \text{target} * \text{input}`, if ``False`` then loss is</span>
<span class="sd">            :math:`\text{input} - \text{target} * \log(\text{input}+\text{eps})`. Default: ``True``</span>
<span class="sd">        full: whether to compute full loss, i. e. to add the Stirling</span>
<span class="sd">            approximation term. Default: ``False``</span>
<span class="sd">            :math:`\text{target} * \log(\text{target}) - \text{target} + 0.5 * \log(2 * \pi * \text{target})`.</span>
<span class="sd">        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,</span>
<span class="sd">            the losses are averaged over each loss element in the batch. Note that for</span>
<span class="sd">            some losses, there multiple elements per sample. If the field :attr:`size_average`</span>
<span class="sd">            is set to ``False``, the losses are instead summed for each minibatch. Ignored</span>
<span class="sd">            when reduce is ``False``. Default: ``True``</span>
<span class="sd">        eps (float, optional): Small value to avoid evaluation of :math:`\log(0)` when</span>
<span class="sd">            :attr:`log_input`\ =\ ``False``. Default: 1e-8</span>
<span class="sd">        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the</span>
<span class="sd">            losses are averaged or summed over observations for each minibatch depending</span>
<span class="sd">            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per</span>
<span class="sd">            batch element instead and ignores :attr:`size_average`. Default: ``True``</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the sum of the output will be divided by the number of</span>
<span class="sd">            elements in the output, ``&#39;sum&#39;``: the output will be summed. Note: :attr:`size_average`</span>
<span class="sd">            and :attr:`reduce` are in the process of being deprecated, and in the meantime,</span>
<span class="sd">            specifying either of those two args will override :attr:`reduction`. Default: ``&#39;mean&#39;``</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">poisson_nll_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">log_input</span><span class="o">=</span><span class="n">log_input</span><span class="p">,</span>
            <span class="n">full</span><span class="o">=</span><span class="n">full</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">reduction</span> <span class="o">!=</span> <span class="s2">&quot;none&quot;</span> <span class="ow">and</span> <span class="n">reduction</span> <span class="o">!=</span> <span class="s2">&quot;mean&quot;</span> <span class="ow">and</span> <span class="n">reduction</span> <span class="o">!=</span> <span class="s2">&quot;sum&quot;</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="nb">input</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="n">reduction</span> <span class="o">+</span> <span class="s2">&quot; is not a valid value for reduction&quot;</span><span class="p">)</span>

    <span class="n">ret</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">poisson_nll_loss</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">log_input</span><span class="p">,</span> <span class="n">full</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<div class="viewcode-block" id="gaussian_nll_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.gaussian_nll_loss.html#torch.nn.functional.gaussian_nll_loss">[docs]</a><span class="k">def</span> <span class="nf">gaussian_nll_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">var</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">full</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">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-6</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Gaussian negative log likelihood loss.</span>

<span class="sd">    See :class:`~torch.nn.GaussianNLLLoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: expectation of the Gaussian distribution.</span>
<span class="sd">        target: sample from the Gaussian distribution.</span>
<span class="sd">        var: tensor of positive variance(s), one for each of the expectations</span>
<span class="sd">            in the input (heteroscedastic), or a single one (homoscedastic).</span>
<span class="sd">        full (bool, optional): include the constant term in the loss calculation. Default: ``False``.</span>
<span class="sd">        eps (float, optional): value added to var, for stability. Default: 1e-6.</span>
<span class="sd">        reduction (str, optional): specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the output is the average of all batch member losses,</span>
<span class="sd">            ``&#39;sum&#39;``: the output is the sum of all batch member losses.</span>
<span class="sd">            Default: ``&#39;mean&#39;``.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">var</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">gaussian_nll_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">var</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">var</span><span class="p">,</span>
            <span class="n">full</span><span class="o">=</span><span class="n">full</span><span class="p">,</span>
            <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="c1"># Check var size</span>
    <span class="c1"># If var.size == input.size, the case is heteroscedastic and no further checks are needed.</span>
    <span class="c1"># Otherwise:</span>
    <span class="k">if</span> <span class="n">var</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">!=</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">():</span>

        <span class="c1"># If var is one dimension short of input, but the sizes match otherwise, then this is a homoscedastic case.</span>
        <span class="c1"># e.g. input.size = (10, 2, 3), var.size = (10, 2)</span>
        <span class="c1"># -&gt; unsqueeze var so that var.shape = (10, 2, 1)</span>
        <span class="c1"># this is done so that broadcasting can happen in the loss calculation</span>
        <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="n">var</span><span class="o">.</span><span class="n">size</span><span class="p">():</span>
            <span class="n">var</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="n">var</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>

        <span class="c1"># This checks if the sizes match up to the final dimension, and the final dimension of var is of size 1.</span>
        <span class="c1"># This is also a homoscedastic case.</span>
        <span class="c1"># e.g. input.size = (10, 2, 3), var.size = (10, 2, 1)</span>
        <span class="k">elif</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">==</span> <span class="n">var</span><span class="o">.</span><span class="n">size</span><span class="p">()[:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="ow">and</span> <span class="n">var</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>  <span class="c1"># Heteroscedastic case</span>
            <span class="k">pass</span>

        <span class="c1"># If none of the above pass, then the size of var is incorrect.</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;var is of incorrect size&quot;</span><span class="p">)</span>

    <span class="c1"># Check validity of reduction mode</span>
    <span class="k">if</span> <span class="n">reduction</span> <span class="o">!=</span> <span class="s1">&#39;none&#39;</span> <span class="ow">and</span> <span class="n">reduction</span> <span class="o">!=</span> <span class="s1">&#39;mean&#39;</span> <span class="ow">and</span> <span class="n">reduction</span> <span class="o">!=</span> <span class="s1">&#39;sum&#39;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="n">reduction</span> <span class="o">+</span> <span class="s2">&quot; is not valid&quot;</span><span class="p">)</span>

    <span class="c1"># Entries of var must be non-negative</span>
    <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">var</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;var has negative entry/entries&quot;</span><span class="p">)</span>

    <span class="c1"># Clamp for stability</span>
    <span class="n">var</span> <span class="o">=</span> <span class="n">var</span><span class="o">.</span><span class="n">clone</span><span class="p">()</span>
    <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">():</span>
        <span class="n">var</span><span class="o">.</span><span class="n">clamp_</span><span class="p">(</span><span class="nb">min</span><span class="o">=</span><span class="n">eps</span><span class="p">)</span>

    <span class="c1"># Calculate the loss</span>
    <span class="n">loss</span> <span class="o">=</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">var</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="nb">input</span> <span class="o">-</span> <span class="n">target</span><span class="p">)</span><span class="o">**</span><span class="mi">2</span> <span class="o">/</span> <span class="n">var</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">full</span><span class="p">:</span>
        <span class="n">loss</span> <span class="o">+=</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="n">math</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="mi">2</span> <span class="o">*</span> <span class="n">math</span><span class="o">.</span><span class="n">pi</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s1">&#39;mean&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">loss</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
    <span class="k">elif</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s1">&#39;sum&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">loss</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">loss</span></div>


<div class="viewcode-block" id="kl_div"><a class="viewcode-back" href="../../../generated/torch.nn.functional.kl_div.html#torch.nn.functional.kl_div">[docs]</a><span class="k">def</span> <span class="nf">kl_div</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
    <span class="n">log_target</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;The `Kullback-Leibler divergence Loss</span>
<span class="sd">    &lt;https://en.wikipedia.org/wiki/Kullback-Leibler_divergence&gt;`__</span>

<span class="sd">    See :class:`~torch.nn.KLDivLoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: Tensor of arbitrary shape in log-probabilities.</span>
<span class="sd">        target: Tensor of the same shape as input. See :attr:`log_target` for</span>
<span class="sd">            the target&#39;s interpretation.</span>
<span class="sd">        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,</span>
<span class="sd">            the losses are averaged over each loss element in the batch. Note that for</span>
<span class="sd">            some losses, there multiple elements per sample. If the field :attr:`size_average`</span>
<span class="sd">            is set to ``False``, the losses are instead summed for each minibatch. Ignored</span>
<span class="sd">            when reduce is ``False``. Default: ``True``</span>
<span class="sd">        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the</span>
<span class="sd">            losses are averaged or summed over observations for each minibatch depending</span>
<span class="sd">            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per</span>
<span class="sd">            batch element instead and ignores :attr:`size_average`. Default: ``True``</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;batchmean&#39;`` | ``&#39;sum&#39;`` | ``&#39;mean&#39;``.</span>
<span class="sd">            ``&#39;none&#39;``: no reduction will be applied</span>
<span class="sd">            ``&#39;batchmean&#39;``: the sum of the output will be divided by the batchsize</span>
<span class="sd">            ``&#39;sum&#39;``: the output will be summed</span>
<span class="sd">            ``&#39;mean&#39;``: the output will be divided by the number of elements in the output</span>
<span class="sd">            Default: ``&#39;mean&#39;``</span>
<span class="sd">        log_target (bool): A flag indicating whether ``target`` is passed in the log space.</span>
<span class="sd">            It is recommended to pass certain distributions (like ``softmax``)</span>
<span class="sd">            in the log space to avoid numerical issues caused by explicit ``log``.</span>
<span class="sd">            Default: ``False``</span>

<span class="sd">    .. note::</span>
<span class="sd">        :attr:`size_average` and :attr:`reduce` are in the process of being deprecated,</span>
<span class="sd">        and in the meantime, specifying either of those two args will override :attr:`reduction`.</span>

<span class="sd">    .. note::</span>
<span class="sd">        :attr:`reduction` = ``&#39;mean&#39;`` doesn&#39;t return the true kl divergence value, please use</span>
<span class="sd">        :attr:`reduction` = ``&#39;batchmean&#39;`` which aligns with KL math definition.</span>
<span class="sd">        In the next major release, ``&#39;mean&#39;`` will be changed to be the same as &#39;batchmean&#39;.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">kl_div</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
            <span class="n">log_target</span><span class="o">=</span><span class="n">log_target</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;mean&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="s2">&quot;reduction: &#39;mean&#39; divides the total loss by both the batch size and the support size.&quot;</span>
                <span class="s2">&quot;&#39;batchmean&#39; divides only by the batch size, and aligns with the KL div math definition.&quot;</span>
                <span class="s2">&quot;&#39;mean&#39; will be changed to behave the same as &#39;batchmean&#39; in the next major release.&quot;</span>
            <span class="p">)</span>

        <span class="c1"># special case for batchmean</span>
        <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;batchmean&quot;</span><span class="p">:</span>
            <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="s2">&quot;sum&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>

    <span class="n">reduced</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">kl_div</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">,</span> <span class="n">log_target</span><span class="o">=</span><span class="n">log_target</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;batchmean&quot;</span> <span class="ow">and</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">reduced</span> <span class="o">=</span> <span class="n">reduced</span> <span class="o">/</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()[</span><span class="mi">0</span><span class="p">]</span>

    <span class="k">return</span> <span class="n">reduced</span></div>


<div class="viewcode-block" id="cross_entropy"><a class="viewcode-back" href="../../../generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy">[docs]</a><span class="k">def</span> <span class="nf">cross_entropy</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">ignore_index</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="o">-</span><span class="mi">100</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
    <span class="n">label_smoothing</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;This criterion computes the cross entropy loss between input logits and target.</span>

<span class="sd">    See :class:`~torch.nn.CrossEntropyLoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor) : Predicted unnormalized logits;</span>
<span class="sd">            see Shape section below for supported shapes.</span>
<span class="sd">        target (Tensor) : Ground truth class indices or class probabilities;</span>
<span class="sd">            see Shape section below for supported shapes.</span>
<span class="sd">        weight (Tensor, optional): a manual rescaling weight given to each</span>
<span class="sd">            class. If given, has to be a Tensor of size `C`</span>
<span class="sd">        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,</span>
<span class="sd">            the losses are averaged over each loss element in the batch. Note that for</span>
<span class="sd">            some losses, there multiple elements per sample. If the field :attr:`size_average`</span>
<span class="sd">            is set to ``False``, the losses are instead summed for each minibatch. Ignored</span>
<span class="sd">            when reduce is ``False``. Default: ``True``</span>
<span class="sd">        ignore_index (int, optional): Specifies a target value that is ignored</span>
<span class="sd">            and does not contribute to the input gradient. When :attr:`size_average` is</span>
<span class="sd">            ``True``, the loss is averaged over non-ignored targets. Note that</span>
<span class="sd">            :attr:`ignore_index` is only applicable when the target contains class indices.</span>
<span class="sd">            Default: -100</span>
<span class="sd">        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the</span>
<span class="sd">            losses are averaged or summed over observations for each minibatch depending</span>
<span class="sd">            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per</span>
<span class="sd">            batch element instead and ignores :attr:`size_average`. Default: ``True``</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the sum of the output will be divided by the number of</span>
<span class="sd">            elements in the output, ``&#39;sum&#39;``: the output will be summed. Note: :attr:`size_average`</span>
<span class="sd">            and :attr:`reduce` are in the process of being deprecated, and in the meantime,</span>
<span class="sd">            specifying either of those two args will override :attr:`reduction`. Default: ``&#39;mean&#39;``</span>
<span class="sd">        label_smoothing (float, optional): A float in [0.0, 1.0]. Specifies the amount</span>
<span class="sd">            of smoothing when computing the loss, where 0.0 means no smoothing. The targets</span>
<span class="sd">            become a mixture of the original ground truth and a uniform distribution as described in</span>
<span class="sd">            `Rethinking the Inception Architecture for Computer Vision &lt;https://arxiv.org/abs/1512.00567&gt;`__. Default: :math:`0.0`.</span>

<span class="sd">    Shape:</span>
<span class="sd">        - Input: Shape :math:`(C)`, :math:`(N, C)` or :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \geq 1`</span>
<span class="sd">          in the case of `K`-dimensional loss.</span>
<span class="sd">        - Target: If containing class indices, shape :math:`()`, :math:`(N)` or :math:`(N, d_1, d_2, ..., d_K)` with</span>
<span class="sd">          :math:`K \geq 1` in the case of K-dimensional loss where each value should be between :math:`[0, C)`.</span>
<span class="sd">          If containing class probabilities, same shape as the input and each value should be between :math:`[0, 1]`.</span>

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

<span class="sd">        .. math::</span>
<span class="sd">            \begin{aligned}</span>
<span class="sd">                C ={} &amp; \text{number of classes} \\</span>
<span class="sd">                N ={} &amp; \text{batch size} \\</span>
<span class="sd">            \end{aligned}</span>

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

<span class="sd">        &gt;&gt;&gt; # Example of target with class indices</span>
<span class="sd">        &gt;&gt;&gt; input = torch.randn(3, 5, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; target = torch.randint(5, (3,), dtype=torch.int64)</span>
<span class="sd">        &gt;&gt;&gt; loss = F.cross_entropy(input, target)</span>
<span class="sd">        &gt;&gt;&gt; loss.backward()</span>
<span class="sd">        &gt;&gt;&gt;</span>
<span class="sd">        &gt;&gt;&gt; # Example of target with class probabilities</span>
<span class="sd">        &gt;&gt;&gt; input = torch.randn(3, 5, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; target = torch.randn(3, 5).softmax(dim=1)</span>
<span class="sd">        &gt;&gt;&gt; loss = F.cross_entropy(input, target)</span>
<span class="sd">        &gt;&gt;&gt; loss.backward()</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">cross_entropy</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">ignore_index</span><span class="o">=</span><span class="n">ignore_index</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
            <span class="n">label_smoothing</span><span class="o">=</span><span class="n">label_smoothing</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">cross_entropy_loss</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">),</span> <span class="n">ignore_index</span><span class="p">,</span> <span class="n">label_smoothing</span><span class="p">)</span></div>


<div class="viewcode-block" id="binary_cross_entropy"><a class="viewcode-back" href="../../../generated/torch.nn.functional.binary_cross_entropy.html#torch.nn.functional.binary_cross_entropy">[docs]</a><span class="k">def</span> <span class="nf">binary_cross_entropy</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Function that measures the Binary Cross Entropy between the target and input</span>
<span class="sd">    probabilities.</span>

<span class="sd">    See :class:`~torch.nn.BCELoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: Tensor of arbitrary shape as probabilities.</span>
<span class="sd">        target: Tensor of the same shape as input with values between 0 and 1.</span>
<span class="sd">        weight (Tensor, optional): a manual rescaling weight</span>
<span class="sd">                if provided it&#39;s repeated to match input tensor shape</span>
<span class="sd">        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,</span>
<span class="sd">            the losses are averaged over each loss element in the batch. Note that for</span>
<span class="sd">            some losses, there multiple elements per sample. If the field :attr:`size_average`</span>
<span class="sd">            is set to ``False``, the losses are instead summed for each minibatch. Ignored</span>
<span class="sd">            when reduce is ``False``. Default: ``True``</span>
<span class="sd">        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the</span>
<span class="sd">            losses are averaged or summed over observations for each minibatch depending</span>
<span class="sd">            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per</span>
<span class="sd">            batch element instead and ignores :attr:`size_average`. Default: ``True``</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the sum of the output will be divided by the number of</span>
<span class="sd">            elements in the output, ``&#39;sum&#39;``: the output will be summed. Note: :attr:`size_average`</span>
<span class="sd">            and :attr:`reduce` are in the process of being deprecated, and in the meantime,</span>
<span class="sd">            specifying either of those two args will override :attr:`reduction`. Default: ``&#39;mean&#39;``</span>

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

<span class="sd">        &gt;&gt;&gt; input = torch.randn(3, 2, requires_grad=True)</span>
<span class="sd">        &gt;&gt;&gt; target = torch.rand(3, 2, requires_grad=False)</span>
<span class="sd">        &gt;&gt;&gt; loss = F.binary_cross_entropy(torch.sigmoid(input), target)</span>
<span class="sd">        &gt;&gt;&gt; loss.backward()</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">binary_cross_entropy</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">!=</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;Using a target size (</span><span class="si">{}</span><span class="s2">) that is different to the input size (</span><span class="si">{}</span><span class="s2">) is deprecated. &quot;</span>
            <span class="s2">&quot;Please ensure they have the same size.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">new_size</span> <span class="o">=</span> <span class="n">_infer_size</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="n">weight</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
        <span class="n">weight</span> <span class="o">=</span> <span class="n">weight</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">new_size</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">binary_cross_entropy</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="binary_cross_entropy_with_logits"><a class="viewcode-back" href="../../../generated/torch.nn.functional.binary_cross_entropy_with_logits.html#torch.nn.functional.binary_cross_entropy_with_logits">[docs]</a><span class="k">def</span> <span class="nf">binary_cross_entropy_with_logits</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
    <span class="n">pos_weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Function that measures Binary Cross Entropy between target and input</span>
<span class="sd">    logits.</span>

<span class="sd">    See :class:`~torch.nn.BCEWithLogitsLoss` for details.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: Tensor of arbitrary shape as unnormalized scores (often referred to as logits).</span>
<span class="sd">        target: Tensor of the same shape as input with values between 0 and 1</span>
<span class="sd">        weight (Tensor, optional): a manual rescaling weight</span>
<span class="sd">            if provided it&#39;s repeated to match input tensor shape</span>
<span class="sd">        size_average (bool, optional): Deprecated (see :attr:`reduction`). By default,</span>
<span class="sd">            the losses are averaged over each loss element in the batch. Note that for</span>
<span class="sd">            some losses, there multiple elements per sample. If the field :attr:`size_average`</span>
<span class="sd">            is set to ``False``, the losses are instead summed for each minibatch. Ignored</span>
<span class="sd">            when reduce is ``False``. Default: ``True``</span>
<span class="sd">        reduce (bool, optional): Deprecated (see :attr:`reduction`). By default, the</span>
<span class="sd">            losses are averaged or summed over observations for each minibatch depending</span>
<span class="sd">            on :attr:`size_average`. When :attr:`reduce` is ``False``, returns a loss per</span>
<span class="sd">            batch element instead and ignores :attr:`size_average`. Default: ``True``</span>
<span class="sd">        reduction (str, optional): Specifies the reduction to apply to the output:</span>
<span class="sd">            ``&#39;none&#39;`` | ``&#39;mean&#39;`` | ``&#39;sum&#39;``. ``&#39;none&#39;``: no reduction will be applied,</span>
<span class="sd">            ``&#39;mean&#39;``: the sum of the output will be divided by the number of</span>
<span class="sd">            elements in the output, ``&#39;sum&#39;``: the output will be summed. Note: :attr:`size_average`</span>
<span class="sd">            and :attr:`reduce` are in the process of being deprecated, and in the meantime,</span>
<span class="sd">            specifying either of those two args will override :attr:`reduction`. Default: ``&#39;mean&#39;``</span>
<span class="sd">        pos_weight (Tensor, optional): a weight of positive examples.</span>
<span class="sd">                Must be a vector with length equal to the number of classes.</span>

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

<span class="sd">         &gt;&gt;&gt; input = torch.randn(3, requires_grad=True)</span>
<span class="sd">         &gt;&gt;&gt; target = torch.empty(3).random_(2)</span>
<span class="sd">         &gt;&gt;&gt; loss = F.binary_cross_entropy_with_logits(input, target)</span>
<span class="sd">         &gt;&gt;&gt; loss.backward()</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">pos_weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">binary_cross_entropy_with_logits</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">pos_weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
            <span class="n">pos_weight</span><span class="o">=</span><span class="n">pos_weight</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">==</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Target size (</span><span class="si">{}</span><span class="s2">) must be the same as input size (</span><span class="si">{}</span><span class="s2">)&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()))</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">binary_cross_entropy_with_logits</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">pos_weight</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="smooth_l1_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.smooth_l1_loss.html#torch.nn.functional.smooth_l1_loss">[docs]</a><span class="k">def</span> <span class="nf">smooth_l1_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
    <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Function that uses a squared term if the absolute</span>
<span class="sd">    element-wise error falls below beta and an L1 term otherwise.</span>

<span class="sd">    See :class:`~torch.nn.SmoothL1Loss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">smooth_l1_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
            <span class="n">beta</span><span class="o">=</span><span class="n">beta</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">==</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</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;Using a target size (</span><span class="si">{}</span><span class="s2">) that is different to the input size (</span><span class="si">{}</span><span class="s2">). &quot;</span>
            <span class="s2">&quot;This will likely lead to incorrect results due to broadcasting. &quot;</span>
            <span class="s2">&quot;Please ensure they have the same size.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()),</span>
            <span class="n">stacklevel</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>

    <span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">broadcast_tensors</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">smooth_l1_loss</span><span class="p">(</span><span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">),</span> <span class="n">beta</span><span class="p">)</span></div>


<div class="viewcode-block" id="huber_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.huber_loss.html#torch.nn.functional.huber_loss">[docs]</a><span class="k">def</span> <span class="nf">huber_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;mean&#39;</span><span class="p">,</span>
    <span class="n">delta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Function that uses a squared term if the absolute</span>
<span class="sd">    element-wise error falls below delta and a delta-scaled L1 term otherwise.</span>

<span class="sd">    See :class:`~torch.nn.HuberLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">huber_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
            <span class="n">delta</span><span class="o">=</span><span class="n">delta</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">==</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</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;Using a target size (</span><span class="si">{}</span><span class="s2">) that is different to the input size (</span><span class="si">{}</span><span class="s2">). &quot;</span>
                      <span class="s2">&quot;This will likely lead to incorrect results due to broadcasting. &quot;</span>
                      <span class="s2">&quot;Please ensure they have the same size.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()),</span>
                      <span class="n">stacklevel</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>

    <span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">broadcast_tensors</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">huber_loss</span><span class="p">(</span><span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">),</span> <span class="n">delta</span><span class="p">)</span></div>


<div class="viewcode-block" id="l1_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.l1_loss.html#torch.nn.functional.l1_loss">[docs]</a><span class="k">def</span> <span class="nf">l1_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;l1_loss(input, target, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    Function that takes the mean element-wise absolute value difference.</span>

<span class="sd">    See :class:`~torch.nn.L1Loss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">l1_loss</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">==</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</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;Using a target size (</span><span class="si">{}</span><span class="s2">) that is different to the input size (</span><span class="si">{}</span><span class="s2">). &quot;</span>
            <span class="s2">&quot;This will likely lead to incorrect results due to broadcasting. &quot;</span>
            <span class="s2">&quot;Please ensure they have the same size.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()),</span>
            <span class="n">stacklevel</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>

    <span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">broadcast_tensors</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">l1_loss</span><span class="p">(</span><span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">))</span></div>


<div class="viewcode-block" id="mse_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.mse_loss.html#torch.nn.functional.mse_loss">[docs]</a><span class="k">def</span> <span class="nf">mse_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;mse_loss(input, target, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    Measures the element-wise mean squared error.</span>

<span class="sd">    See :class:`~torch.nn.MSELoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">mse_loss</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">()</span> <span class="o">==</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</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;Using a target size (</span><span class="si">{}</span><span class="s2">) that is different to the input size (</span><span class="si">{}</span><span class="s2">). &quot;</span>
            <span class="s2">&quot;This will likely lead to incorrect results due to broadcasting. &quot;</span>
            <span class="s2">&quot;Please ensure they have the same size.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">()),</span>
            <span class="n">stacklevel</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>

    <span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">broadcast_tensors</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">mse_loss</span><span class="p">(</span><span class="n">expanded_input</span><span class="p">,</span> <span class="n">expanded_target</span><span class="p">,</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">))</span></div>


<div class="viewcode-block" id="margin_ranking_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.margin_ranking_loss.html#torch.nn.functional.margin_ranking_loss">[docs]</a><span class="k">def</span> <span class="nf">margin_ranking_loss</span><span class="p">(</span>
    <span class="n">input1</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">input2</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;margin_ranking_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.MarginRankingLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="n">input1</span><span class="p">,</span> <span class="n">input2</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">margin_ranking_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="n">input1</span><span class="p">,</span> <span class="n">input2</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="n">input1</span><span class="p">,</span>
            <span class="n">input2</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">margin</span><span class="o">=</span><span class="n">margin</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">if</span> <span class="p">(</span><span class="n">input1</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="n">input2</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="ow">or</span> <span class="n">input1</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="n">target</span><span class="o">.</span><span class="n">dim</span><span class="p">()):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
            <span class="p">(</span>
                <span class="s2">&quot;margin_ranking_loss : All input tensors should have same dimension but got sizes: &quot;</span>
                <span class="s2">&quot;input1: </span><span class="si">{}</span><span class="s2">, input2: </span><span class="si">{}</span><span class="s2">, target: </span><span class="si">{}</span><span class="s2"> &quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">input1</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="n">input2</span><span class="o">.</span><span class="n">size</span><span class="p">(),</span> <span class="n">target</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>
            <span class="p">)</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">margin_ranking_loss</span><span class="p">(</span><span class="n">input1</span><span class="p">,</span> <span class="n">input2</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">margin</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="hinge_embedding_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.hinge_embedding_loss.html#torch.nn.functional.hinge_embedding_loss">[docs]</a><span class="k">def</span> <span class="nf">hinge_embedding_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;hinge_embedding_loss(input, target, margin=1.0, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.HingeEmbeddingLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">hinge_embedding_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">margin</span><span class="o">=</span><span class="n">margin</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">hinge_embedding_loss</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">margin</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="multilabel_margin_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.multilabel_margin_loss.html#torch.nn.functional.multilabel_margin_loss">[docs]</a><span class="k">def</span> <span class="nf">multilabel_margin_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;multilabel_margin_loss(input, target, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.MultiLabelMarginLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">multilabel_margin_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">multilabel_margin_loss</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="soft_margin_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.soft_margin_loss.html#torch.nn.functional.soft_margin_loss">[docs]</a><span class="k">def</span> <span class="nf">soft_margin_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;soft_margin_loss(input, target, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.SoftMarginLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">soft_margin_loss</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">soft_margin_loss</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="multilabel_soft_margin_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.multilabel_soft_margin_loss.html#torch.nn.functional.multilabel_soft_margin_loss">[docs]</a><span class="k">def</span> <span class="nf">multilabel_soft_margin_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;multilabel_soft_margin_loss(input, target, weight=None, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.MultiLabelSoftMarginLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">multilabel_soft_margin_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_string</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>

    <span class="n">loss</span> <span class="o">=</span> <span class="o">-</span><span class="p">(</span><span class="n">target</span> <span class="o">*</span> <span class="n">logsigmoid</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">target</span><span class="p">)</span> <span class="o">*</span> <span class="n">logsigmoid</span><span class="p">(</span><span class="o">-</span><span class="nb">input</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span> <span class="o">*</span> <span class="n">weight</span>

    <span class="n">class_dim</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">-</span> <span class="mi">1</span>
    <span class="n">C</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">class_dim</span><span class="p">)</span>
    <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="n">class_dim</span><span class="p">)</span> <span class="o">/</span> <span class="n">C</span>  <span class="c1"># only return N loss values</span>

    <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;none&quot;</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">loss</span>
    <span class="k">elif</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;mean&quot;</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">loss</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
    <span class="k">elif</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;sum&quot;</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="n">loss</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ret</span> <span class="o">=</span> <span class="nb">input</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="n">reduction</span> <span class="o">+</span> <span class="s2">&quot; is not valid&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">ret</span></div>


<div class="viewcode-block" id="cosine_embedding_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.cosine_embedding_loss.html#torch.nn.functional.cosine_embedding_loss">[docs]</a><span class="k">def</span> <span class="nf">cosine_embedding_loss</span><span class="p">(</span>
    <span class="n">input1</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">input2</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;cosine_embedding_loss(input1, input2, target, margin=0, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.CosineEmbeddingLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="n">input1</span><span class="p">,</span> <span class="n">input2</span><span class="p">,</span> <span class="n">target</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">cosine_embedding_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="n">input1</span><span class="p">,</span> <span class="n">input2</span><span class="p">,</span> <span class="n">target</span><span class="p">),</span>
            <span class="n">input1</span><span class="p">,</span>
            <span class="n">input2</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">margin</span><span class="o">=</span><span class="n">margin</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">cosine_embedding_loss</span><span class="p">(</span><span class="n">input1</span><span class="p">,</span> <span class="n">input2</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">margin</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="multi_margin_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.multi_margin_loss.html#torch.nn.functional.multi_margin_loss">[docs]</a><span class="k">def</span> <span class="nf">multi_margin_loss</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">target</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">p</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
    <span class="n">weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;multi_margin_loss(input, target, p=1, margin=1, weight=None, size_average=None, reduce=None, reduction=&#39;mean&#39;) -&gt; Tensor</span>

<span class="sd">    See :class:`~torch.nn.MultiMarginLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">multi_margin_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">weight</span><span class="p">),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">target</span><span class="p">,</span>
            <span class="n">p</span><span class="o">=</span><span class="n">p</span><span class="p">,</span>
            <span class="n">margin</span><span class="o">=</span><span class="n">margin</span><span class="p">,</span>
            <span class="n">weight</span><span class="o">=</span><span class="n">weight</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">p</span> <span class="o">!=</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">p</span> <span class="o">!=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;only p == 1 and p == 2 supported&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">weight</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;weight must be one-dimensional&quot;</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">multi_margin_loss</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">target</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">margin</span><span class="p">,</span> <span class="n">weight</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<span class="n">pixel_shuffle</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">pixel_shuffle</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">pixel_shuffle(input, upscale_factor) -&gt; Tensor</span>

<span class="sd">Rearranges elements in a tensor of shape :math:`(*, C \times r^2, H, W)` to a</span>
<span class="sd">tensor of shape :math:`(*, C, H \times r, W \times r)`, where r is the :attr:`upscale_factor`.</span>

<span class="sd">See :class:`~torch.nn.PixelShuffle` for details.</span>

<span class="sd">Args:</span>
<span class="sd">    input (Tensor): the input tensor</span>
<span class="sd">    upscale_factor (int): factor to increase spatial resolution by</span>

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

<span class="sd">    &gt;&gt;&gt; input = torch.randn(1, 9, 4, 4)</span>
<span class="sd">    &gt;&gt;&gt; output = torch.nn.functional.pixel_shuffle(input, 3)</span>
<span class="sd">    &gt;&gt;&gt; print(output.size())</span>
<span class="sd">    torch.Size([1, 1, 12, 12])</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">pixel_unshuffle</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">pixel_unshuffle</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">pixel_unshuffle(input, downscale_factor) -&gt; Tensor</span>

<span class="sd">Reverses the :class:`~torch.nn.PixelShuffle` operation by rearranging elements in a</span>
<span class="sd">tensor of shape :math:`(*, C, H \times r, W \times r)` to a tensor of shape</span>
<span class="sd">:math:`(*, C \times r^2, H, W)`, where r is the :attr:`downscale_factor`.</span>

<span class="sd">See :class:`~torch.nn.PixelUnshuffle` for details.</span>

<span class="sd">Args:</span>
<span class="sd">    input (Tensor): the input tensor</span>
<span class="sd">    downscale_factor (int): factor to increase spatial resolution by</span>

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

<span class="sd">    &gt;&gt;&gt; input = torch.randn(1, 1, 12, 12)</span>
<span class="sd">    &gt;&gt;&gt; output = torch.nn.functional.pixel_unshuffle(input, 3)</span>
<span class="sd">    &gt;&gt;&gt; print(output.size())</span>
<span class="sd">    torch.Size([1, 9, 4, 4])</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">channel_shuffle</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">channel_shuffle</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">channel_shuffle(input, groups) -&gt; Tensor</span>

<span class="sd">Divide the channels in a tensor of shape :math:`(*, C , H, W)`</span>
<span class="sd">into g groups and rearrange them as :math:`(*, C \frac g, g, H, W)`,</span>
<span class="sd">while keeping the original tensor shape.</span>

<span class="sd">See :class:`~torch.nn.ChannelShuffle` for details.</span>

<span class="sd">Args:</span>
<span class="sd">    input (Tensor): the input tensor</span>
<span class="sd">    groups (int): number of groups to divide channels in and rearrange.</span>

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

<span class="sd">    &gt;&gt;&gt; input = torch.randn(1, 4, 2, 2)</span>
<span class="sd">    &gt;&gt;&gt; print(input)</span>
<span class="sd">    [[[[1, 2],</span>
<span class="sd">       [3, 4]],</span>
<span class="sd">      [[5, 6],</span>
<span class="sd">       [7, 8]],</span>
<span class="sd">      [[9, 10],</span>
<span class="sd">       [11, 12]],</span>
<span class="sd">      [[13, 14],</span>
<span class="sd">       [15, 16]],</span>
<span class="sd">     ]]</span>
<span class="sd">    &gt;&gt;&gt; output = torch.nn.functional.channel_shuffle(input, 2)</span>
<span class="sd">    &gt;&gt;&gt; print(output)</span>
<span class="sd">    [[[[1, 2],</span>
<span class="sd">       [3, 4]],</span>
<span class="sd">      [[9, 10],</span>
<span class="sd">       [11, 12]],</span>
<span class="sd">      [[5, 6],</span>
<span class="sd">       [7, 8]],</span>
<span class="sd">      [[13, 14],</span>
<span class="sd">       [15, 16]],</span>
<span class="sd">     ]]</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">native_channel_shuffle</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">native_channel_shuffle</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">native_channel_shuffle(input, groups) -&gt; Tensor</span>

<span class="sd">Native kernel level implementation of the `channel_shuffle`.</span>
<span class="sd">This function might become private in future releases, use with caution.</span>

<span class="sd">Divide the channels in a tensor of shape :math:`(*, C , H, W)`</span>
<span class="sd">into g groups and rearrange them as :math:`(*, C \frac g, g, H, W)`,</span>
<span class="sd">while keeping the original tensor shape.</span>

<span class="sd">See :class:`~torch.nn.ChannelShuffle` for details.</span>

<span class="sd">Args:</span>
<span class="sd">    input (Tensor): the input tensor</span>
<span class="sd">    groups (int): number of groups to divide channels in and rearrange.</span>

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

<span class="sd">    &gt;&gt;&gt; input = torch.randn(1, 4, 2, 2)</span>
<span class="sd">    &gt;&gt;&gt; print(input)</span>
<span class="sd">    [[[[1, 2],</span>
<span class="sd">       [3, 4]],</span>
<span class="sd">      [[5, 6],</span>
<span class="sd">       [7, 8]],</span>
<span class="sd">      [[9, 10],</span>
<span class="sd">       [11, 12]],</span>
<span class="sd">      [[13, 14],</span>
<span class="sd">       [15, 16]],</span>
<span class="sd">     ]]</span>
<span class="sd">    &gt;&gt;&gt; output = torch.nn.functional.native_channel_shuffle(input, 2)</span>
<span class="sd">    &gt;&gt;&gt; print(output)</span>
<span class="sd">    [[[[1, 2],</span>
<span class="sd">       [3, 4]],</span>
<span class="sd">      [[9, 10],</span>
<span class="sd">       [11, 12]],</span>
<span class="sd">      [[5, 6],</span>
<span class="sd">       [7, 8]],</span>
<span class="sd">      [[13, 14],</span>
<span class="sd">       [15, 16]],</span>
<span class="sd">     ]]</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;nearest&quot;</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;nearest&quot;</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<div class="viewcode-block" id="upsample"><a class="viewcode-back" href="../../../generated/torch.nn.functional.upsample.html#torch.nn.functional.upsample">[docs]</a><span class="k">def</span> <span class="nf">upsample</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">scale_factor</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;nearest&quot;</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Upsamples the input to either the given :attr:`size` or the given</span>
<span class="sd">    :attr:`scale_factor`</span>

<span class="sd">    .. warning::</span>
<span class="sd">        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.</span>
<span class="sd">        This is equivalent with ``nn.functional.interpolate(...)``.</span>

<span class="sd">    Note:</span>
<span class="sd">        {backward_reproducibility_note}</span>

<span class="sd">    The algorithm used for upsampling is determined by :attr:`mode`.</span>

<span class="sd">    Currently temporal, spatial and volumetric upsampling are supported, i.e.</span>
<span class="sd">    expected inputs are 3-D, 4-D or 5-D in shape.</span>

<span class="sd">    The input dimensions are interpreted in the form:</span>
<span class="sd">    `mini-batch x channels x [optional depth] x [optional height] x width`.</span>

<span class="sd">    The modes available for upsampling are: `nearest`, `linear` (3D-only),</span>
<span class="sd">    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only)</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): the input tensor</span>
<span class="sd">        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):</span>
<span class="sd">            output spatial size.</span>
<span class="sd">        scale_factor (float or Tuple[float]): multiplier for spatial size. Has to match input size if it is a tuple.</span>
<span class="sd">        mode (str): algorithm used for upsampling:</span>
<span class="sd">            ``&#39;nearest&#39;`` | ``&#39;linear&#39;`` | ``&#39;bilinear&#39;`` | ``&#39;bicubic&#39;`` |</span>
<span class="sd">            ``&#39;trilinear&#39;``. Default: ``&#39;nearest&#39;``</span>
<span class="sd">        align_corners (bool, optional): Geometrically, we consider the pixels of the</span>
<span class="sd">            input and output as squares rather than points.</span>
<span class="sd">            If set to ``True``, the input and output tensors are aligned by the</span>
<span class="sd">            center points of their corner pixels, preserving the values at the corner pixels.</span>
<span class="sd">            If set to ``False``, the input and output tensors are aligned by the corner</span>
<span class="sd">            points of their corner pixels, and the interpolation uses edge value padding</span>
<span class="sd">            for out-of-boundary values, making this operation *independent* of input size</span>
<span class="sd">            when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`</span>
<span class="sd">            is ``&#39;linear&#39;``, ``&#39;bilinear&#39;``, ``&#39;bicubic&#39;`` or ``&#39;trilinear&#39;``.</span>
<span class="sd">            Default: ``False``</span>

<span class="sd">    .. note::</span>
<span class="sd">        With ``mode=&#39;bicubic&#39;``, it&#39;s possible to cause overshoot, in other words it can produce</span>
<span class="sd">        negative values or values greater than 255 for images.</span>
<span class="sd">        Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot</span>
<span class="sd">        when displaying the image.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        With ``align_corners = True``, the linearly interpolating modes</span>
<span class="sd">        (`linear`, `bilinear`, and `trilinear`) don&#39;t proportionally align the</span>
<span class="sd">        output and input pixels, and thus the output values can depend on the</span>
<span class="sd">        input size. This was the default behavior for these modes up to version</span>
<span class="sd">        0.3.1. Since then, the default behavior is ``align_corners = False``.</span>
<span class="sd">        See :class:`~torch.nn.Upsample` for concrete examples on how this</span>
<span class="sd">        affects the outputs.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;nn.functional.upsample is deprecated. Use nn.functional.interpolate instead.&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">scale_factor</span><span class="p">,</span> <span class="n">mode</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">)</span></div>


<span class="k">if</span> <span class="n">upsample</span><span class="o">.</span><span class="vm">__doc__</span><span class="p">:</span>
    <span class="n">upsample</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">upsample</span><span class="o">.</span><span class="vm">__doc__</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;nearest&#39;</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">recompute_scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">antialias</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;nearest&#39;</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">recompute_scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">antialias</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;nearest&#39;</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">recompute_scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">antialias</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">interpolate</span><span class="p">(</span>  <span class="c1"># noqa: F811</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</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="p">,</span>
    <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;nearest&quot;</span><span class="p">,</span>
    <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">recompute_scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">antialias</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>

<div class="viewcode-block" id="interpolate"><a class="viewcode-back" href="../../../generated/torch.nn.functional.interpolate.html#torch.nn.functional.interpolate">[docs]</a><span class="k">def</span> <span class="nf">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;nearest&#39;</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">recompute_scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">antialias</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="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Down/up samples the input to either the given :attr:`size` or the given</span>
<span class="sd">    :attr:`scale_factor`</span>

<span class="sd">    The algorithm used for interpolation is determined by :attr:`mode`.</span>

<span class="sd">    Currently temporal, spatial and volumetric sampling are supported, i.e.</span>
<span class="sd">    expected inputs are 3-D, 4-D or 5-D in shape.</span>

<span class="sd">    The input dimensions are interpreted in the form:</span>
<span class="sd">    `mini-batch x channels x [optional depth] x [optional height] x width`.</span>

<span class="sd">    The modes available for resizing are: `nearest`, `linear` (3D-only),</span>
<span class="sd">    `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `area`, `nearest-exact`</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): the input tensor</span>
<span class="sd">        size (int or Tuple[int] or Tuple[int, int] or Tuple[int, int, int]):</span>
<span class="sd">            output spatial size.</span>
<span class="sd">        scale_factor (float or Tuple[float]): multiplier for spatial size. If `scale_factor` is a tuple,</span>
<span class="sd">            its length has to match the number of spatial dimensions; `input.dim() - 2`.</span>
<span class="sd">        mode (str): algorithm used for upsampling:</span>
<span class="sd">            ``&#39;nearest&#39;`` | ``&#39;linear&#39;`` | ``&#39;bilinear&#39;`` | ``&#39;bicubic&#39;`` |</span>
<span class="sd">            ``&#39;trilinear&#39;`` | ``&#39;area&#39;`` | ``&#39;nearest-exact&#39;``. Default: ``&#39;nearest&#39;``</span>
<span class="sd">        align_corners (bool, optional): Geometrically, we consider the pixels of the</span>
<span class="sd">            input and output as squares rather than points.</span>
<span class="sd">            If set to ``True``, the input and output tensors are aligned by the</span>
<span class="sd">            center points of their corner pixels, preserving the values at the corner pixels.</span>
<span class="sd">            If set to ``False``, the input and output tensors are aligned by the corner</span>
<span class="sd">            points of their corner pixels, and the interpolation uses edge value padding</span>
<span class="sd">            for out-of-boundary values, making this operation *independent* of input size</span>
<span class="sd">            when :attr:`scale_factor` is kept the same. This only has an effect when :attr:`mode`</span>
<span class="sd">            is ``&#39;linear&#39;``, ``&#39;bilinear&#39;``, ``&#39;bicubic&#39;`` or ``&#39;trilinear&#39;``.</span>
<span class="sd">            Default: ``False``</span>
<span class="sd">        recompute_scale_factor (bool, optional): recompute the scale_factor for use in the</span>
<span class="sd">            interpolation calculation. If `recompute_scale_factor` is ``True``, then</span>
<span class="sd">            `scale_factor` must be passed in and `scale_factor` is used to compute the</span>
<span class="sd">            output `size`. The computed output `size` will be used to infer new scales for</span>
<span class="sd">            the interpolation. Note that when `scale_factor` is floating-point, it may differ</span>
<span class="sd">            from the recomputed `scale_factor` due to rounding and precision issues.</span>
<span class="sd">            If `recompute_scale_factor` is ``False``, then `size` or `scale_factor` will</span>
<span class="sd">            be used directly for interpolation. Default: ``None``.</span>
<span class="sd">        antialias (bool, optional): flag to apply anti-aliasing. Default: ``False``. Using anti-alias</span>
<span class="sd">            option together with ``align_corners=False``, interpolation result would match Pillow</span>
<span class="sd">            result for downsampling operation. Supported modes: ``&#39;bilinear&#39;``, ``&#39;bicubic&#39;``.</span>

<span class="sd">    .. note::</span>
<span class="sd">        With ``mode=&#39;bicubic&#39;``, it&#39;s possible to cause overshoot, in other words it can produce</span>
<span class="sd">        negative values or values greater than 255 for images.</span>
<span class="sd">        Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot</span>
<span class="sd">        when displaying the image.</span>

<span class="sd">    .. note::</span>
<span class="sd">        Mode ``mode=&#39;nearest-exact&#39;`` matches Scikit-Image and PIL nearest neighbours interpolation</span>
<span class="sd">        algorithms and fixes known issues with ``mode=&#39;nearest&#39;``. This mode is introduced to keep</span>
<span class="sd">        backward compatibility.</span>
<span class="sd">        Mode ``mode=&#39;nearest&#39;`` matches buggy OpenCV&#39;s ``INTER_NEAREST`` interpolation algorithm.</span>

<span class="sd">    Note:</span>
<span class="sd">        {backward_reproducibility_note}</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">interpolate</span><span class="p">,</span>
            <span class="p">(</span><span class="nb">input</span><span class="p">,),</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">size</span><span class="o">=</span><span class="n">size</span><span class="p">,</span>
            <span class="n">scale_factor</span><span class="o">=</span><span class="n">scale_factor</span><span class="p">,</span>
            <span class="n">mode</span><span class="o">=</span><span class="n">mode</span><span class="p">,</span>
            <span class="n">align_corners</span><span class="o">=</span><span class="n">align_corners</span><span class="p">,</span>
            <span class="n">recompute_scale_factor</span><span class="o">=</span><span class="n">recompute_scale_factor</span><span class="p">,</span>
            <span class="n">antialias</span><span class="o">=</span><span class="n">antialias</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">mode</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">&quot;nearest&quot;</span><span class="p">,</span> <span class="s2">&quot;area&quot;</span><span class="p">,</span> <span class="s2">&quot;nearest-exact&quot;</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;align_corners option can only be set with the &quot;</span>
                <span class="s2">&quot;interpolating modes: linear | bilinear | bicubic | trilinear&quot;</span>
            <span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">align_corners</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="n">dim</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">-</span> <span class="mi">2</span>  <span class="c1"># Number of spatial dimensions.</span>

    <span class="c1"># Process size and scale_factor.  Validate that exactly one is set.</span>
    <span class="c1"># Validate its length if it is a list, or expand it if it is a scalar.</span>
    <span class="c1"># After this block, exactly one of output_size and scale_factors will</span>
    <span class="c1"># be non-None, and it will be a list (or tuple).</span>
    <span class="k">if</span> <span class="n">size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">scale_factor</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;only one of size or scale_factor should be defined&quot;</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">scale_factor</span> <span class="ow">is</span> <span class="kc">None</span>
        <span class="n">scale_factors</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">size</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">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">size</span><span class="p">)</span> <span class="o">!=</span> <span class="n">dim</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s2">&quot;Input and output must have the same number of spatial dimensions, but got &quot;</span>
                    <span class="sa">f</span><span class="s2">&quot;input with spatial dimensions of </span><span class="si">{</span><span class="nb">list</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:])</span><span class="si">}</span><span class="s2"> and output size of </span><span class="si">{</span><span class="n">size</span><span class="si">}</span><span class="s2">. &quot;</span>
                    <span class="s2">&quot;Please provide input tensor in (N, C, d1, d2, ...,dK) format and &quot;</span>
                    <span class="s2">&quot;output size in (o1, o2, ...,oK) format.&quot;</span>

                <span class="p">)</span>
            <span class="n">output_size</span> <span class="o">=</span> <span class="n">size</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">output_size</span> <span class="o">=</span> <span class="p">[</span><span class="n">size</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="n">dim</span><span class="p">)]</span>
    <span class="k">elif</span> <span class="n">scale_factor</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">size</span> <span class="ow">is</span> <span class="kc">None</span>
        <span class="n">output_size</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">scale_factor</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">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">scale_factor</span><span class="p">)</span> <span class="o">!=</span> <span class="n">dim</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s2">&quot;Input and scale_factor must have the same number of spatial dimensions, but &quot;</span>
                    <span class="sa">f</span><span class="s2">&quot;got input with spatial dimensions of </span><span class="si">{</span><span class="nb">list</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">:])</span><span class="si">}</span><span class="s2"> and &quot;</span>
                    <span class="sa">f</span><span class="s2">&quot;scale_factor of shape </span><span class="si">{</span><span class="n">scale_factor</span><span class="si">}</span><span class="s2">. &quot;</span>
                    <span class="s2">&quot;Please provide input tensor in (N, C, d1, d2, ...,dK) format and &quot;</span>
                    <span class="s2">&quot;scale_factor in (s1, s2, ...,sK) format.&quot;</span>
                <span class="p">)</span>
            <span class="n">scale_factors</span> <span class="o">=</span> <span class="n">scale_factor</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">scale_factors</span> <span class="o">=</span> <span class="p">[</span><span class="n">scale_factor</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="n">dim</span><span class="p">)]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;either size or scale_factor should be defined&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">recompute_scale_factor</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">recompute_scale_factor</span> <span class="ow">and</span> <span class="n">size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;recompute_scale_factor is not meaningful with an explicit size.&quot;</span><span class="p">)</span>

    <span class="c1"># &quot;area&quot; mode always requires an explicit size rather than scale factor.</span>
    <span class="c1"># Re-use the recompute_scale_factor code path.</span>
    <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;area&quot;</span> <span class="ow">and</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">recompute_scale_factor</span> <span class="o">=</span> <span class="kc">True</span>

    <span class="k">if</span> <span class="n">recompute_scale_factor</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">recompute_scale_factor</span><span class="p">:</span>
        <span class="c1"># We compute output_size here, then un-set scale_factors.</span>
        <span class="c1"># The C++ code will recompute it based on the (integer) output size.</span>
        <span class="k">assert</span> <span class="n">scale_factors</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">()</span> <span class="ow">and</span> <span class="n">torch</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="c1"># make scale_factor a tensor in tracing so constant doesn&#39;t get baked in</span>
            <span class="n">output_size</span> <span class="o">=</span> <span class="p">[</span>
                <span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">floor</span><span class="p">((</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">float</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">scale_factors</span><span class="p">[</span><span class="n">i</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="o">.</span><span class="n">float</span><span class="p">()))</span>
                <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
            <span class="p">]</span>
        <span class="k">elif</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">():</span>
            <span class="n">output_size</span> <span class="o">=</span> <span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">math</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="mi">2</span><span class="p">))</span> <span class="o">*</span> <span class="n">scale_factors</span><span class="p">[</span><span class="n">i</span><span class="p">]))</span>
                           <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">dim</span><span class="p">)]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">output_size</span> <span class="o">=</span> <span class="p">[</span>
                <span class="n">_sym_int</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="n">i</span> <span class="o">+</span> <span class="mi">2</span><span class="p">)</span> <span class="o">*</span> <span class="n">scale_factors</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
                <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">dim</span><span class="p">)</span>
            <span class="p">]</span>
        <span class="n">scale_factors</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="k">if</span> <span class="n">antialias</span> <span class="ow">and</span> <span class="ow">not</span> <span class="p">(</span><span class="n">mode</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">&quot;bilinear&quot;</span><span class="p">,</span> <span class="s2">&quot;bicubic&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">input</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Anti-alias option is only supported for bilinear and bicubic modes&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_nearest1d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_nearest2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">5</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_nearest3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest-exact&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">_upsample_nearest_exact1d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest-exact&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">_upsample_nearest_exact2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">5</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest-exact&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">_upsample_nearest_exact3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;area&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">return</span> <span class="n">adaptive_avg_pool1d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;area&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">return</span> <span class="n">adaptive_avg_pool2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">5</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;area&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">output_size</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">return</span> <span class="n">adaptive_avg_pool3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;linear&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_linear1d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;bilinear&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="n">antialias</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">_upsample_bilinear2d_aa</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_bilinear2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">5</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;trilinear&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_trilinear3d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;bicubic&quot;</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
        <span class="k">if</span> <span class="n">antialias</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">_upsample_bicubic2d_aa</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">upsample_bicubic2d</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">,</span> <span class="n">scale_factors</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;bilinear&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Got 3D input, but bilinear mode needs 4D input&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;trilinear&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Got 3D input, but trilinear mode needs 5D input&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;linear&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Got 4D input, but linear mode needs 3D input&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;trilinear&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Got 4D input, but trilinear mode needs 5D input&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">5</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;linear&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Got 5D input, but linear mode needs 3D input&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">5</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;bilinear&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span><span class="s2">&quot;Got 5D input, but bilinear mode needs 4D input&quot;</span><span class="p">)</span>

    <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span>
        <span class="s2">&quot;Input Error: Only 3D, 4D and 5D input Tensors supported&quot;</span>
        <span class="s2">&quot; (got </span><span class="si">{}</span><span class="s2">D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact&quot;</span>
        <span class="s2">&quot; (got </span><span class="si">{}</span><span class="s2">)&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">dim</span><span class="p">(),</span> <span class="n">mode</span><span class="p">)</span>
    <span class="p">)</span></div>


<span class="k">if</span> <span class="n">interpolate</span><span class="o">.</span><span class="vm">__doc__</span><span class="p">:</span>
    <span class="n">interpolate</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">interpolate</span><span class="o">.</span><span class="vm">__doc__</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample_nearest</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample_nearest</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<div class="viewcode-block" id="upsample_nearest"><a class="viewcode-back" href="../../../generated/torch.nn.functional.upsample_nearest.html#torch.nn.functional.upsample_nearest">[docs]</a><span class="k">def</span> <span class="nf">upsample_nearest</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">scale_factor</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Upsamples the input, using nearest neighbours&#39; pixel values.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.</span>
<span class="sd">        This is equivalent with ``nn.functional.interpolate(..., mode=&#39;nearest&#39;)``.</span>

<span class="sd">    Currently spatial and volumetric upsampling are supported (i.e. expected</span>
<span class="sd">    inputs are 4 or 5 dimensional).</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input</span>
<span class="sd">        size (int or Tuple[int, int] or Tuple[int, int, int]): output spatia</span>
<span class="sd">            size.</span>
<span class="sd">        scale_factor (int): multiplier for spatial size. Has to be an integer.</span>

<span class="sd">    Note:</span>
<span class="sd">        {backward_reproducibility_note}</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># DeprecationWarning is ignored by default</span>
    <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;nn.functional.upsample_nearest is deprecated. Use nn.functional.interpolate instead.&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">scale_factor</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;nearest&quot;</span><span class="p">)</span></div>


<span class="k">if</span> <span class="n">upsample_nearest</span><span class="o">.</span><span class="vm">__doc__</span><span class="p">:</span>
    <span class="n">upsample_nearest</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">upsample_nearest</span><span class="o">.</span><span class="vm">__doc__</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample_bilinear</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample_bilinear</span><span class="p">(</span>  <span class="c1"># noqa: F811</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">float</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample_bilinear</span><span class="p">(</span>  <span class="c1"># noqa: F811</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<span class="nd">@_overload</span>  <span class="c1"># noqa: F811</span>
<span class="k">def</span> <span class="nf">upsample_bilinear</span><span class="p">(</span>  <span class="c1"># noqa: F811</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</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="p">,</span> <span class="n">scale_factor</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">float</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>  <span class="c1"># noqa: F811</span>
    <span class="k">pass</span>


<div class="viewcode-block" id="upsample_bilinear"><a class="viewcode-back" href="../../../generated/torch.nn.functional.upsample_bilinear.html#torch.nn.functional.upsample_bilinear">[docs]</a><span class="k">def</span> <span class="nf">upsample_bilinear</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">scale_factor</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Upsamples the input, using bilinear upsampling.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        This function is deprecated in favor of :func:`torch.nn.functional.interpolate`.</span>
<span class="sd">        This is equivalent with</span>
<span class="sd">        ``nn.functional.interpolate(..., mode=&#39;bilinear&#39;, align_corners=True)``.</span>

<span class="sd">    Expected inputs are spatial (4 dimensional). Use `upsample_trilinear` fo</span>
<span class="sd">    volumetric (5 dimensional) inputs.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input</span>
<span class="sd">        size (int or Tuple[int, int]): output spatial size.</span>
<span class="sd">        scale_factor (int or Tuple[int, int]): multiplier for spatial size</span>

<span class="sd">    Note:</span>
<span class="sd">        {backward_reproducibility_note}</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># DeprecationWarning is ignored by default</span>
    <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;nn.functional.upsample_bilinear is deprecated. Use nn.functional.interpolate instead.&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">interpolate</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">scale_factor</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="s2">&quot;bilinear&quot;</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span></div>


<span class="k">if</span> <span class="n">upsample_bilinear</span><span class="o">.</span><span class="vm">__doc__</span><span class="p">:</span>
    <span class="n">upsample_bilinear</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">upsample_bilinear</span><span class="o">.</span><span class="vm">__doc__</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>

<span class="n">GRID_SAMPLE_INTERPOLATION_MODES</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;bilinear&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span>
    <span class="s2">&quot;nearest&quot;</span><span class="p">:</span> <span class="mi">1</span><span class="p">,</span>
    <span class="s2">&quot;bicubic&quot;</span><span class="p">:</span> <span class="mi">2</span><span class="p">,</span>
<span class="p">}</span>

<span class="n">GRID_SAMPLE_PADDING_MODES</span> <span class="o">=</span> <span class="p">{</span>
    <span class="s2">&quot;zeros&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span>
    <span class="s2">&quot;border&quot;</span><span class="p">:</span> <span class="mi">1</span><span class="p">,</span>
    <span class="s2">&quot;reflection&quot;</span><span class="p">:</span> <span class="mi">2</span><span class="p">,</span>
<span class="p">}</span>


<div class="viewcode-block" id="grid_sample"><a class="viewcode-back" href="../../../generated/torch.nn.functional.grid_sample.html#torch.nn.functional.grid_sample">[docs]</a><span class="k">def</span> <span class="nf">grid_sample</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">grid</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;bilinear&quot;</span><span class="p">,</span>
    <span class="n">padding_mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;zeros&quot;</span><span class="p">,</span>
    <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Given an :attr:`input` and a flow-field :attr:`grid`, computes the</span>
<span class="sd">    ``output`` using :attr:`input` values and pixel locations from :attr:`grid`.</span>

<span class="sd">    Currently, only spatial (4-D) and volumetric (5-D) :attr:`input` are</span>
<span class="sd">    supported.</span>

<span class="sd">    In the spatial (4-D) case, for :attr:`input` with shape</span>
<span class="sd">    :math:`(N, C, H_\text{in}, W_\text{in})` and :attr:`grid` with shape</span>
<span class="sd">    :math:`(N, H_\text{out}, W_\text{out}, 2)`, the output will have shape</span>
<span class="sd">    :math:`(N, C, H_\text{out}, W_\text{out})`.</span>

<span class="sd">    For each output location ``output[n, :, h, w]``, the size-2 vector</span>
<span class="sd">    ``grid[n, h, w]`` specifies :attr:`input` pixel locations ``x`` and ``y``,</span>
<span class="sd">    which are used to interpolate the output value ``output[n, :, h, w]``.</span>
<span class="sd">    In the case of 5D inputs, ``grid[n, d, h, w]`` specifies the</span>
<span class="sd">    ``x``, ``y``, ``z`` pixel locations for interpolating</span>
<span class="sd">    ``output[n, :, d, h, w]``. :attr:`mode` argument specifies ``nearest`` or</span>
<span class="sd">    ``bilinear`` interpolation method to sample the input pixels.</span>

<span class="sd">    :attr:`grid` specifies the sampling pixel locations normalized by the</span>
<span class="sd">    :attr:`input` spatial dimensions. Therefore, it should have most values in</span>
<span class="sd">    the range of ``[-1, 1]``. For example, values ``x = -1, y = -1`` is the</span>
<span class="sd">    left-top pixel of :attr:`input`, and values  ``x = 1, y = 1`` is the</span>
<span class="sd">    right-bottom pixel of :attr:`input`.</span>

<span class="sd">    If :attr:`grid` has values outside the range of ``[-1, 1]``, the corresponding</span>
<span class="sd">    outputs are handled as defined by :attr:`padding_mode`. Options are</span>

<span class="sd">        * ``padding_mode=&quot;zeros&quot;``: use ``0`` for out-of-bound grid locations,</span>
<span class="sd">        * ``padding_mode=&quot;border&quot;``: use border values for out-of-bound grid locations,</span>
<span class="sd">        * ``padding_mode=&quot;reflection&quot;``: use values at locations reflected by</span>
<span class="sd">          the border for out-of-bound grid locations. For location far away</span>
<span class="sd">          from the border, it will keep being reflected until becoming in bound,</span>
<span class="sd">          e.g., (normalized) pixel location ``x = -3.5`` reflects by border ``-1``</span>
<span class="sd">          and becomes ``x&#39; = 1.5``, then reflects by border ``1`` and becomes</span>
<span class="sd">          ``x&#39;&#39; = -0.5``.</span>

<span class="sd">    Note:</span>
<span class="sd">        This function is often used in conjunction with :func:`affine_grid`</span>
<span class="sd">        to build `Spatial Transformer Networks`_ .</span>

<span class="sd">    Note:</span>
<span class="sd">        When using the CUDA backend, this operation may induce nondeterministic</span>
<span class="sd">        behaviour in its backward pass that is not easily switched off.</span>
<span class="sd">        Please see the notes on :doc:`/notes/randomness` for background.</span>

<span class="sd">    Note:</span>
<span class="sd">        NaN values in :attr:`grid` would be interpreted as ``-1``.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): input of shape :math:`(N, C, H_\text{in}, W_\text{in})` (4-D case)</span>
<span class="sd">                        or :math:`(N, C, D_\text{in}, H_\text{in}, W_\text{in})` (5-D case)</span>
<span class="sd">        grid (Tensor): flow-field of shape :math:`(N, H_\text{out}, W_\text{out}, 2)` (4-D case)</span>
<span class="sd">                       or :math:`(N, D_\text{out}, H_\text{out}, W_\text{out}, 3)` (5-D case)</span>
<span class="sd">        mode (str): interpolation mode to calculate output values</span>
<span class="sd">            ``&#39;bilinear&#39;`` | ``&#39;nearest&#39;`` | ``&#39;bicubic&#39;``. Default: ``&#39;bilinear&#39;``</span>
<span class="sd">            Note: ``mode=&#39;bicubic&#39;`` supports only 4-D input.</span>
<span class="sd">            When ``mode=&#39;bilinear&#39;`` and the input is 5-D, the interpolation mode</span>
<span class="sd">            used internally will actually be trilinear. However, when the input is 4-D,</span>
<span class="sd">            the interpolation mode will legitimately be bilinear.</span>
<span class="sd">        padding_mode (str): padding mode for outside grid values</span>
<span class="sd">            ``&#39;zeros&#39;`` | ``&#39;border&#39;`` | ``&#39;reflection&#39;``. Default: ``&#39;zeros&#39;``</span>
<span class="sd">        align_corners (bool, optional): Geometrically, we consider the pixels of the</span>
<span class="sd">            input  as squares rather than points.</span>
<span class="sd">            If set to ``True``, the extrema (``-1`` and ``1``) are considered as referring</span>
<span class="sd">            to the center points of the input&#39;s corner pixels. If set to ``False``, they</span>
<span class="sd">            are instead considered as referring to the corner points of the input&#39;s corner</span>
<span class="sd">            pixels, making the sampling more resolution agnostic.</span>
<span class="sd">            This option parallels the ``align_corners`` option in</span>
<span class="sd">            :func:`interpolate`, and so whichever option is used here</span>
<span class="sd">            should also be used there to resize the input image before grid sampling.</span>
<span class="sd">            Default: ``False``</span>

<span class="sd">    Returns:</span>
<span class="sd">        output (Tensor): output Tensor</span>

<span class="sd">    .. _`Spatial Transformer Networks`:</span>
<span class="sd">        https://arxiv.org/abs/1506.02025</span>

<span class="sd">    .. warning::</span>
<span class="sd">        When ``align_corners = True``, the grid positions depend on the pixel</span>
<span class="sd">        size relative to the input image size, and so the locations sampled by</span>
<span class="sd">        :func:`grid_sample` will differ for the same input given at different</span>
<span class="sd">        resolutions (that is, after being upsampled or downsampled).</span>
<span class="sd">        The default behavior up to version 1.2.0 was ``align_corners = True``.</span>
<span class="sd">        Since then, the default behavior has been changed to ``align_corners = False``,</span>
<span class="sd">        in order to bring it in line with the default for :func:`interpolate`.</span>

<span class="sd">    .. note::</span>
<span class="sd">        ``mode=&#39;bicubic&#39;`` is implemented using the `cubic convolution algorithm`_ with :math:`\alpha=-0.75`.</span>
<span class="sd">        The constant :math:`\alpha` might be different from packages to packages.</span>
<span class="sd">        For example, `PIL`_ and `OpenCV`_ use -0.5 and -0.75 respectively.</span>
<span class="sd">        This algorithm may &quot;overshoot&quot; the range of values it&#39;s interpolating.</span>
<span class="sd">        For example, it may produce negative values or values greater than 255 when interpolating input in [0, 255].</span>
<span class="sd">        Clamp the results with :func: `torch.clamp` to ensure they are within the valid range.</span>
<span class="sd">    .. _`cubic convolution algorithm`: https://en.wikipedia.org/wiki/Bicubic_interpolation</span>
<span class="sd">    .. _`PIL`: https://github.com/python-pillow/Pillow/blob/4634eafe3c695a014267eefdce830b4a825beed7/src/libImaging/Resample.c#L51</span>
<span class="sd">    .. _`OpenCV`: https://github.com/opencv/opencv/blob/f345ed564a06178670750bad59526cfa4033be55/modules/imgproc/src/resize.cpp#L908</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">grid</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">grid_sample</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">grid</span><span class="p">),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">grid</span><span class="p">,</span> <span class="n">mode</span><span class="o">=</span><span class="n">mode</span><span class="p">,</span> <span class="n">padding_mode</span><span class="o">=</span><span class="n">padding_mode</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="n">align_corners</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">mode</span> <span class="o">!=</span> <span class="s2">&quot;bilinear&quot;</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">!=</span> <span class="s2">&quot;nearest&quot;</span> <span class="ow">and</span> <span class="n">mode</span> <span class="o">!=</span> <span class="s2">&quot;bicubic&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;nn.functional.grid_sample(): expected mode to be &quot;</span>
            <span class="s2">&quot;&#39;bilinear&#39;, &#39;nearest&#39; or &#39;bicubic&#39;, but got: &#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">mode</span><span class="p">)</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">padding_mode</span> <span class="o">!=</span> <span class="s2">&quot;zeros&quot;</span> <span class="ow">and</span> <span class="n">padding_mode</span> <span class="o">!=</span> <span class="s2">&quot;border&quot;</span> <span class="ow">and</span> <span class="n">padding_mode</span> <span class="o">!=</span> <span class="s2">&quot;reflection&quot;</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
            <span class="s2">&quot;nn.functional.grid_sample(): expected padding_mode &quot;</span>
            <span class="s2">&quot;to be &#39;zeros&#39;, &#39;border&#39;, or &#39;reflection&#39;, &quot;</span>
            <span class="s2">&quot;but got: &#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">padding_mode</span><span class="p">)</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;bilinear&quot;</span><span class="p">:</span>
        <span class="n">mode_enum</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">elif</span> <span class="n">mode</span> <span class="o">==</span> <span class="s2">&quot;nearest&quot;</span><span class="p">:</span>
        <span class="n">mode_enum</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="k">else</span><span class="p">:</span>  <span class="c1"># mode == &#39;bicubic&#39;</span>
        <span class="n">mode_enum</span> <span class="o">=</span> <span class="mi">2</span>

    <span class="k">if</span> <span class="n">padding_mode</span> <span class="o">==</span> <span class="s2">&quot;zeros&quot;</span><span class="p">:</span>
        <span class="n">padding_mode_enum</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">elif</span> <span class="n">padding_mode</span> <span class="o">==</span> <span class="s2">&quot;border&quot;</span><span class="p">:</span>
        <span class="n">padding_mode_enum</span> <span class="o">=</span> <span class="mi">1</span>
    <span class="k">else</span><span class="p">:</span>  <span class="c1"># padding_mode == &#39;reflection&#39;</span>
        <span class="n">padding_mode_enum</span> <span class="o">=</span> <span class="mi">2</span>

    <span class="k">if</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="kc">None</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;Default grid_sample and affine_grid behavior has changed &quot;</span>
            <span class="s2">&quot;to align_corners=False since 1.3.0. Please specify &quot;</span>
            <span class="s2">&quot;align_corners=True if the old behavior is desired. &quot;</span>
            <span class="s2">&quot;See the documentation of grid_sample for details.&quot;</span>
        <span class="p">)</span>
        <span class="n">align_corners</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">grid_sampler</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">grid</span><span class="p">,</span> <span class="n">mode_enum</span><span class="p">,</span> <span class="n">padding_mode_enum</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">)</span></div>


<div class="viewcode-block" id="affine_grid"><a class="viewcode-back" href="../../../generated/torch.nn.functional.affine_grid.html#torch.nn.functional.affine_grid">[docs]</a><span class="k">def</span> <span class="nf">affine_grid</span><span class="p">(</span><span class="n">theta</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">size</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">align_corners</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Generates a 2D or 3D flow field (sampling grid), given a batch of</span>
<span class="sd">    affine matrices :attr:`theta`.</span>

<span class="sd">    .. note::</span>
<span class="sd">        This function is often used in conjunction with :func:`grid_sample`</span>
<span class="sd">        to build `Spatial Transformer Networks`_ .</span>

<span class="sd">    Args:</span>
<span class="sd">        theta (Tensor): input batch of affine matrices with shape</span>
<span class="sd">            (:math:`N \times 2 \times 3`) for 2D or</span>
<span class="sd">            (:math:`N \times 3 \times 4`) for 3D</span>
<span class="sd">        size (torch.Size): the target output image size.</span>
<span class="sd">            (:math:`N \times C \times H \times W` for 2D or</span>
<span class="sd">            :math:`N \times C \times D \times H \times W` for 3D)</span>
<span class="sd">            Example: torch.Size((32, 3, 24, 24))</span>
<span class="sd">        align_corners (bool, optional): if ``True``, consider ``-1`` and ``1``</span>
<span class="sd">            to refer to the centers of the corner pixels rather than the image corners.</span>
<span class="sd">            Refer to :func:`grid_sample` for a more complete description.</span>
<span class="sd">            A grid generated by :func:`affine_grid` should be passed to :func:`grid_sample`</span>
<span class="sd">            with the same setting for this option.</span>
<span class="sd">            Default: ``False``</span>

<span class="sd">    Returns:</span>
<span class="sd">        output (Tensor): output Tensor of size (:math:`N \times H \times W \times 2`)</span>

<span class="sd">    .. _`Spatial Transformer Networks`:</span>
<span class="sd">        https://arxiv.org/abs/1506.02025</span>

<span class="sd">    .. warning::</span>
<span class="sd">        When ``align_corners = True``, the grid positions depend on the pixel</span>
<span class="sd">        size relative to the input image size, and so the locations sampled by</span>
<span class="sd">        :func:`grid_sample` will differ for the same input given at different</span>
<span class="sd">        resolutions (that is, after being upsampled or downsampled).</span>
<span class="sd">        The default behavior up to version 1.2.0 was ``align_corners = True``.</span>
<span class="sd">        Since then, the default behavior has been changed to ``align_corners = False``,</span>
<span class="sd">        in order to bring it in line with the default for :func:`interpolate`.</span>
<span class="sd">    .. warning::</span>
<span class="sd">        When ``align_corners = True``, 2D affine transforms on 1D data and</span>
<span class="sd">        3D affine transforms on 2D data (that is, when one of the spatial</span>
<span class="sd">        dimensions has unit size) are ill-defined, and not an intended use case.</span>
<span class="sd">        This is not a problem when ``align_corners = False``.</span>
<span class="sd">        Up to version 1.2.0, all grid points along a unit dimension were</span>
<span class="sd">        considered arbitrarily to be at ``-1``.</span>
<span class="sd">        From version 1.3.0, under ``align_corners = True`` all grid points</span>
<span class="sd">        along a unit dimension are considered to be at ``0``</span>
<span class="sd">        (the center of the input image).</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="n">theta</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">affine_grid</span><span class="p">,</span> <span class="p">(</span><span class="n">theta</span><span class="p">,),</span> <span class="n">theta</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">align_corners</span><span class="o">=</span><span class="n">align_corners</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">align_corners</span> <span class="ow">is</span> <span class="kc">None</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;Default grid_sample and affine_grid behavior has changed &quot;</span>
            <span class="s2">&quot;to align_corners=False since 1.3.0. Please specify &quot;</span>
            <span class="s2">&quot;align_corners=True if the old behavior is desired. &quot;</span>
            <span class="s2">&quot;See the documentation of grid_sample for details.&quot;</span>
        <span class="p">)</span>
        <span class="n">align_corners</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="c1"># enforce floating point dtype on theta</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">theta</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Expected theta to have floating point type, but got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">theta</span><span class="o">.</span><span class="n">dtype</span><span class="p">))</span>
    <span class="c1"># check that shapes and sizes match</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">size</span><span class="p">)</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">theta</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="mi">3</span> <span class="ow">or</span> <span class="n">theta</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">2</span> <span class="ow">or</span> <span class="n">theta</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">3</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;Expected a batch of 2D affine matrices of shape Nx2x3 &quot;</span>
                <span class="s2">&quot;for size </span><span class="si">{}</span><span class="s2">. Got </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">theta</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
            <span class="p">)</span>
        <span class="n">spatial_size</span> <span class="o">=</span> <span class="n">size</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>  <span class="c1"># spatial dimension sizes</span>
    <span class="k">elif</span> <span class="nb">len</span><span class="p">(</span><span class="n">size</span><span class="p">)</span> <span class="o">==</span> <span class="mi">5</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">theta</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">!=</span> <span class="mi">3</span> <span class="ow">or</span> <span class="n">theta</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">3</span> <span class="ow">or</span> <span class="n">theta</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">!=</span> <span class="mi">4</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                <span class="s2">&quot;Expected a batch of 3D affine matrices of shape Nx3x4 &quot;</span>
                <span class="s2">&quot;for size </span><span class="si">{}</span><span class="s2">. Got </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">size</span><span class="p">,</span> <span class="n">theta</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
            <span class="p">)</span>
        <span class="n">spatial_size</span> <span class="o">=</span> <span class="n">size</span><span class="p">[</span><span class="o">-</span><span class="mi">3</span><span class="p">:]</span>  <span class="c1"># spatial dimension sizes</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span>
            <span class="s2">&quot;affine_grid only supports 4D and 5D sizes, &quot;</span>
            <span class="s2">&quot;for 2D and 3D affine transforms, respectively. &quot;</span>
            <span class="s2">&quot;Got size </span><span class="si">{}</span><span class="s2">.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">size</span><span class="p">)</span>
        <span class="p">)</span>
    <span class="c1"># check for empty span</span>
    <span class="k">if</span> <span class="n">align_corners</span> <span class="ow">and</span> <span class="nb">min</span><span class="p">(</span><span class="n">spatial_size</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</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;Since version 1.3.0, affine_grid behavior has changed &quot;</span>
            <span class="s2">&quot;for unit-size grids when align_corners=True. &quot;</span>
            <span class="s2">&quot;This is not an intended use case of affine_grid. &quot;</span>
            <span class="s2">&quot;See the documentation of affine_grid for details.&quot;</span>
        <span class="p">)</span>
    <span class="k">elif</span> <span class="nb">min</span><span class="p">(</span><span class="n">size</span><span class="p">)</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Expected non-zero, positive output size. Got </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">size</span><span class="p">))</span>

    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">affine_grid_generator</span><span class="p">(</span><span class="n">theta</span><span class="p">,</span> <span class="n">size</span><span class="p">,</span> <span class="n">align_corners</span><span class="p">)</span></div>


<span class="n">pad</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">pad</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">pad(input, pad, mode=&quot;constant&quot;, value=None) -&gt; Tensor</span>

<span class="sd">Pads tensor.</span>

<span class="sd">Padding size:</span>
<span class="sd">    The padding size by which to pad some dimensions of :attr:`input`</span>
<span class="sd">    are described starting from the last dimension and moving forward.</span>
<span class="sd">    :math:`\left\lfloor\frac{\text{len(pad)}}{2}\right\rfloor` dimensions</span>
<span class="sd">    of ``input`` will be padded.</span>
<span class="sd">    For example, to pad only the last dimension of the input tensor, then</span>
<span class="sd">    :attr:`pad` has the form</span>
<span class="sd">    :math:`(\text{padding\_left}, \text{padding\_right})`;</span>
<span class="sd">    to pad the last 2 dimensions of the input tensor, then use</span>
<span class="sd">    :math:`(\text{padding\_left}, \text{padding\_right},`</span>
<span class="sd">    :math:`\text{padding\_top}, \text{padding\_bottom})`;</span>
<span class="sd">    to pad the last 3 dimensions, use</span>
<span class="sd">    :math:`(\text{padding\_left}, \text{padding\_right},`</span>
<span class="sd">    :math:`\text{padding\_top}, \text{padding\_bottom}`</span>
<span class="sd">    :math:`\text{padding\_front}, \text{padding\_back})`.</span>

<span class="sd">Padding mode:</span>
<span class="sd">    See :class:`torch.nn.ConstantPad2d`, :class:`torch.nn.ReflectionPad2d`, and</span>
<span class="sd">    :class:`torch.nn.ReplicationPad2d` for concrete examples on how each of the</span>
<span class="sd">    padding modes works. Constant padding is implemented for arbitrary dimensions.</span>
<span class="sd">    Replicate and reflection padding are implemented for padding the last 3</span>
<span class="sd">    dimensions of a 4D or 5D input tensor, the last 2 dimensions of a 3D</span>
<span class="sd">    or 4D input tensor, or the last dimension of a 2D or 3D input tensor.</span>

<span class="sd">Note:</span>
<span class="sd">    When using the CUDA backend, this operation may induce nondeterministic</span>
<span class="sd">    behaviour in its backward pass that is not easily switched off.</span>
<span class="sd">    Please see the notes on :doc:`/notes/randomness` for background.</span>

<span class="sd">Args:</span>
<span class="sd">    input (Tensor): N-dimensional tensor</span>
<span class="sd">    pad (tuple): m-elements tuple, where</span>
<span class="sd">        :math:`\frac{m}{2} \leq` input dimensions and :math:`m` is even.</span>
<span class="sd">    mode: ``&#39;constant&#39;``, ``&#39;reflect&#39;``, ``&#39;replicate&#39;`` or ``&#39;circular&#39;``.</span>
<span class="sd">        Default: ``&#39;constant&#39;``</span>
<span class="sd">    value: fill value for ``&#39;constant&#39;`` padding. Default: ``0``</span>

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

<span class="sd">    &gt;&gt;&gt; t4d = torch.empty(3, 3, 4, 2)</span>
<span class="sd">    &gt;&gt;&gt; p1d = (1, 1) # pad last dim by 1 on each side</span>
<span class="sd">    &gt;&gt;&gt; out = F.pad(t4d, p1d, &quot;constant&quot;, 0)  # effectively zero padding</span>
<span class="sd">    &gt;&gt;&gt; print(out.size())</span>
<span class="sd">    torch.Size([3, 3, 4, 4])</span>
<span class="sd">    &gt;&gt;&gt; p2d = (1, 1, 2, 2) # pad last dim by (1, 1) and 2nd to last by (2, 2)</span>
<span class="sd">    &gt;&gt;&gt; out = F.pad(t4d, p2d, &quot;constant&quot;, 0)</span>
<span class="sd">    &gt;&gt;&gt; print(out.size())</span>
<span class="sd">    torch.Size([3, 3, 8, 4])</span>
<span class="sd">    &gt;&gt;&gt; t4d = torch.empty(3, 3, 4, 2)</span>
<span class="sd">    &gt;&gt;&gt; p3d = (0, 1, 2, 1, 3, 3) # pad by (0, 1), (2, 1), and (3, 3)</span>
<span class="sd">    &gt;&gt;&gt; out = F.pad(t4d, p3d, &quot;constant&quot;, 0)</span>
<span class="sd">    &gt;&gt;&gt; print(out.size())</span>
<span class="sd">    torch.Size([3, 9, 7, 3])</span>

<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>
<span class="c1"># TODO: Fix via https://github.com/pytorch/pytorch/issues/75798</span>
<span class="n">pad</span><span class="o">.</span><span class="vm">__module__</span> <span class="o">=</span> <span class="s2">&quot;torch.nn.functional&quot;</span>

<span class="c1"># distance</span>


<span class="n">pairwise_distance</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">pairwise_distance</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">pairwise_distance(x1, x2, p=2.0, eps=1e-6, keepdim=False) -&gt; Tensor</span>

<span class="sd">See :class:`torch.nn.PairwiseDistance` for details</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>


<span class="n">pdist</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">pdist</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">pdist(input, p=2) -&gt; Tensor</span>

<span class="sd">Computes the p-norm distance between every pair of row vectors in the input.</span>
<span class="sd">This is identical to the upper triangular portion, excluding the diagonal, of</span>
<span class="sd">`torch.norm(input[:, None] - input, dim=2, p=p)`. This function will be faster</span>
<span class="sd">if the rows are contiguous.</span>

<span class="sd">If input has shape :math:`N \times M` then the output will have shape</span>
<span class="sd">:math:`\frac{1}{2} N (N - 1)`.</span>

<span class="sd">This function is equivalent to ``scipy.spatial.distance.pdist(input,</span>
<span class="sd">&#39;minkowski&#39;, p=p)`` if :math:`p \in (0, \infty)`. When :math:`p = 0` it is</span>
<span class="sd">equivalent to ``scipy.spatial.distance.pdist(input, &#39;hamming&#39;) * M``.</span>
<span class="sd">When :math:`p = \infty`, the closest scipy function is</span>
<span class="sd">``scipy.spatial.distance.pdist(xn, lambda x, y: np.abs(x - y).max())``.</span>

<span class="sd">Args:</span>
<span class="sd">    input: input tensor of shape :math:`N \times M`.</span>
<span class="sd">    p: p value for the p-norm distance to calculate between each vector pair</span>
<span class="sd">        :math:`\in [0, \infty]`.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="n">cosine_similarity</span> <span class="o">=</span> <span class="n">_add_docstr</span><span class="p">(</span>
    <span class="n">torch</span><span class="o">.</span><span class="n">cosine_similarity</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">cosine_similarity(x1, x2, dim=1, eps=1e-8) -&gt; Tensor</span>

<span class="sd">Returns cosine similarity between ``x1`` and ``x2``, computed along dim. ``x1`` and ``x2`` must be broadcastable</span>
<span class="sd">to a common shape. ``dim`` refers to the dimension in this common shape. Dimension ``dim`` of the output is</span>
<span class="sd">squeezed (see :func:`torch.squeeze`), resulting in the</span>
<span class="sd">output tensor having 1 fewer dimension.</span>

<span class="sd">.. math ::</span>
<span class="sd">    \text{similarity} = \dfrac{x_1 \cdot x_2}{\max(\Vert x_1 \Vert _2 \cdot \Vert x_2 \Vert _2, \epsilon)}</span>

<span class="sd">Supports :ref:`type promotion &lt;type-promotion-doc&gt;`.</span>

<span class="sd">Args:</span>
<span class="sd">    x1 (Tensor): First input.</span>
<span class="sd">    x2 (Tensor): Second input.</span>
<span class="sd">    dim (int, optional): Dimension along which cosine similarity is computed. Default: 1</span>
<span class="sd">    eps (float, optional): Small value to avoid division by zero.</span>
<span class="sd">        Default: 1e-8</span>

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

<span class="sd">    &gt;&gt;&gt; input1 = torch.randn(100, 128)</span>
<span class="sd">    &gt;&gt;&gt; input2 = torch.randn(100, 128)</span>
<span class="sd">    &gt;&gt;&gt; output = F.cosine_similarity(input1, input2)</span>
<span class="sd">    &gt;&gt;&gt; print(output)</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<span class="n">one_hot</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">one_hot</span><span class="p">,</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">one_hot(tensor, num_classes=-1) -&gt; LongTensor</span>

<span class="sd">Takes LongTensor with index values of shape ``(*)`` and returns a tensor</span>
<span class="sd">of shape ``(*, num_classes)`` that have zeros everywhere except where the</span>
<span class="sd">index of last dimension matches the corresponding value of the input tensor,</span>
<span class="sd">in which case it will be 1.</span>

<span class="sd">See also `One-hot on Wikipedia`_ .</span>

<span class="sd">.. _One-hot on Wikipedia:</span>
<span class="sd">    https://en.wikipedia.org/wiki/One-hot</span>

<span class="sd">Arguments:</span>
<span class="sd">    tensor (LongTensor): class values of any shape.</span>
<span class="sd">    num_classes (int):  Total number of classes. If set to -1, the number</span>
<span class="sd">        of classes will be inferred as one greater than the largest class</span>
<span class="sd">        value in the input tensor.</span>

<span class="sd">Returns:</span>
<span class="sd">    LongTensor that has one more dimension with 1 values at the</span>
<span class="sd">    index of last dimension indicated by the input, and 0 everywhere</span>
<span class="sd">    else.</span>

<span class="sd">Examples:</span>
<span class="sd">    &gt;&gt;&gt; F.one_hot(torch.arange(0, 5) % 3)</span>
<span class="sd">    tensor([[1, 0, 0],</span>
<span class="sd">            [0, 1, 0],</span>
<span class="sd">            [0, 0, 1],</span>
<span class="sd">            [1, 0, 0],</span>
<span class="sd">            [0, 1, 0]])</span>
<span class="sd">    &gt;&gt;&gt; F.one_hot(torch.arange(0, 5) % 3, num_classes=5)</span>
<span class="sd">    tensor([[1, 0, 0, 0, 0],</span>
<span class="sd">            [0, 1, 0, 0, 0],</span>
<span class="sd">            [0, 0, 1, 0, 0],</span>
<span class="sd">            [1, 0, 0, 0, 0],</span>
<span class="sd">            [0, 1, 0, 0, 0]])</span>
<span class="sd">    &gt;&gt;&gt; F.one_hot(torch.arange(0, 6).view(3,2) % 3)</span>
<span class="sd">    tensor([[[1, 0, 0],</span>
<span class="sd">             [0, 1, 0]],</span>
<span class="sd">            [[0, 0, 1],</span>
<span class="sd">             [1, 0, 0]],</span>
<span class="sd">            [[0, 1, 0],</span>
<span class="sd">             [0, 0, 1]]])</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">,</span>
<span class="p">)</span>


<div class="viewcode-block" id="triplet_margin_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.triplet_margin_loss.html#torch.nn.functional.triplet_margin_loss">[docs]</a><span class="k">def</span> <span class="nf">triplet_margin_loss</span><span class="p">(</span>
    <span class="n">anchor</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">positive</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">negative</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
    <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">2</span><span class="p">,</span>
    <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-6</span><span class="p">,</span>
    <span class="n">swap</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">size_average</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduce</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">bool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    See :class:`~torch.nn.TripletMarginLoss` for details</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">positive</span><span class="p">,</span> <span class="n">negative</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">triplet_margin_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">positive</span><span class="p">,</span> <span class="n">negative</span><span class="p">),</span>
            <span class="n">anchor</span><span class="p">,</span>
            <span class="n">positive</span><span class="p">,</span>
            <span class="n">negative</span><span class="p">,</span>
            <span class="n">margin</span><span class="o">=</span><span class="n">margin</span><span class="p">,</span>
            <span class="n">p</span><span class="o">=</span><span class="n">p</span><span class="p">,</span>
            <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span>
            <span class="n">swap</span><span class="o">=</span><span class="n">swap</span><span class="p">,</span>
            <span class="n">size_average</span><span class="o">=</span><span class="n">size_average</span><span class="p">,</span>
            <span class="n">reduce</span><span class="o">=</span><span class="n">reduce</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">if</span> <span class="n">size_average</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">reduce</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">legacy_get_enum</span><span class="p">(</span><span class="n">size_average</span><span class="p">,</span> <span class="n">reduce</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">reduction_enum</span> <span class="o">=</span> <span class="n">_Reduction</span><span class="o">.</span><span class="n">get_enum</span><span class="p">(</span><span class="n">reduction</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">triplet_margin_loss</span><span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">positive</span><span class="p">,</span> <span class="n">negative</span><span class="p">,</span> <span class="n">margin</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">swap</span><span class="p">,</span> <span class="n">reduction_enum</span><span class="p">)</span></div>


<div class="viewcode-block" id="triplet_margin_with_distance_loss"><a class="viewcode-back" href="../../../generated/torch.nn.functional.triplet_margin_with_distance_loss.html#torch.nn.functional.triplet_margin_with_distance_loss">[docs]</a><span class="k">def</span> <span class="nf">triplet_margin_with_distance_loss</span><span class="p">(</span>
    <span class="n">anchor</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">positive</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">negative</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="o">*</span><span class="p">,</span>
    <span class="n">distance_function</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Callable</span><span class="p">[[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">],</span> <span class="n">Tensor</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">margin</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
    <span class="n">swap</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">reduction</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;mean&quot;</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    See :class:`~torch.nn.TripletMarginWithDistanceLoss` for details.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">is_scripting</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">NotImplementedError</span><span class="p">(</span>
            <span class="s2">&quot;F.triplet_margin_with_distance_loss does not support JIT scripting: &quot;</span>
            <span class="s2">&quot;functions requiring Callables cannot be scripted.&quot;</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">positive</span><span class="p">,</span> <span class="n">negative</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">triplet_margin_with_distance_loss</span><span class="p">,</span>
            <span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">positive</span><span class="p">,</span> <span class="n">negative</span><span class="p">),</span>
            <span class="n">anchor</span><span class="p">,</span>
            <span class="n">positive</span><span class="p">,</span>
            <span class="n">negative</span><span class="p">,</span>
            <span class="n">distance_function</span><span class="o">=</span><span class="n">distance_function</span><span class="p">,</span>
            <span class="n">margin</span><span class="o">=</span><span class="n">margin</span><span class="p">,</span>
            <span class="n">swap</span><span class="o">=</span><span class="n">swap</span><span class="p">,</span>
            <span class="n">reduction</span><span class="o">=</span><span class="n">reduction</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="c1"># Check validity of reduction mode</span>
    <span class="k">if</span> <span class="n">reduction</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">(</span><span class="s2">&quot;mean&quot;</span><span class="p">,</span> <span class="s2">&quot;sum&quot;</span><span class="p">,</span> <span class="s2">&quot;none&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;</span><span class="si">{</span><span class="n">reduction</span><span class="si">}</span><span class="s2"> is not a valid value for reduction&quot;</span><span class="p">)</span>

    <span class="c1"># Check dimensions</span>
    <span class="n">a_dim</span> <span class="o">=</span> <span class="n">anchor</span><span class="o">.</span><span class="n">ndim</span>
    <span class="n">p_dim</span> <span class="o">=</span> <span class="n">positive</span><span class="o">.</span><span class="n">ndim</span>
    <span class="n">n_dim</span> <span class="o">=</span> <span class="n">negative</span><span class="o">.</span><span class="n">ndim</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="p">(</span><span class="n">a_dim</span> <span class="o">==</span> <span class="n">p_dim</span> <span class="ow">and</span> <span class="n">p_dim</span> <span class="o">==</span> <span class="n">n_dim</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
            <span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The anchor, positive, and negative tensors are expected to have &quot;</span>
             <span class="sa">f</span><span class="s2">&quot;the same number of dimensions, but got: anchor </span><span class="si">{</span><span class="n">a_dim</span><span class="si">}</span><span class="s2">D, &quot;</span>
             <span class="sa">f</span><span class="s2">&quot;positive </span><span class="si">{</span><span class="n">p_dim</span><span class="si">}</span><span class="s2">D, and negative </span><span class="si">{</span><span class="n">n_dim</span><span class="si">}</span><span class="s2">D inputs&quot;</span><span class="p">))</span>

    <span class="c1"># Calculate loss</span>
    <span class="k">if</span> <span class="n">distance_function</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">distance_function</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">pairwise_distance</span>

    <span class="n">dist_pos</span> <span class="o">=</span> <span class="n">distance_function</span><span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">positive</span><span class="p">)</span>
    <span class="n">dist_neg</span> <span class="o">=</span> <span class="n">distance_function</span><span class="p">(</span><span class="n">anchor</span><span class="p">,</span> <span class="n">negative</span><span class="p">)</span>
    <span class="c1"># The distance swap is described in the paper &quot;Learning shallow</span>
    <span class="c1"># convolutional feature descriptors with triplet losses&quot; by V. Balntas, E.</span>
    <span class="c1"># Riba et al.  If True, and if the positive example is closer to the</span>
    <span class="c1"># negative example than the anchor is, swaps the positive example and the</span>
    <span class="c1"># anchor in the loss computation.</span>
    <span class="k">if</span> <span class="n">swap</span><span class="p">:</span>
        <span class="n">dist_swap</span> <span class="o">=</span> <span class="n">distance_function</span><span class="p">(</span><span class="n">positive</span><span class="p">,</span> <span class="n">negative</span><span class="p">)</span>
        <span class="n">dist_neg</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">minimum</span><span class="p">(</span><span class="n">dist_neg</span><span class="p">,</span> <span class="n">dist_swap</span><span class="p">)</span>
    <span class="n">loss</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">clamp_min</span><span class="p">(</span><span class="n">margin</span> <span class="o">+</span> <span class="n">dist_pos</span> <span class="o">-</span> <span class="n">dist_neg</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>

    <span class="c1"># Apply reduction</span>
    <span class="k">if</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;sum&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">reduction</span> <span class="o">==</span> <span class="s2">&quot;mean&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>  <span class="c1"># reduction == &quot;none&quot;</span>
        <span class="k">return</span> <span class="n">loss</span></div>


<div class="viewcode-block" id="normalize"><a class="viewcode-back" href="../../../generated/torch.nn.functional.normalize.html#torch.nn.functional.normalize">[docs]</a><span class="k">def</span> <span class="nf">normalize</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">p</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span> <span class="n">dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span> <span class="n">eps</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-12</span><span class="p">,</span> <span class="n">out</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Performs :math:`L_p` normalization of inputs over specified dimension.</span>

<span class="sd">    For a tensor :attr:`input` of sizes :math:`(n_0, ..., n_{dim}, ..., n_k)`, each</span>
<span class="sd">    :math:`n_{dim}` -element vector :math:`v` along dimension :attr:`dim` is transformed as</span>

<span class="sd">    .. math::</span>
<span class="sd">        v = \frac{v}{\max(\lVert v \rVert_p, \epsilon)}.</span>

<span class="sd">    With the default arguments it uses the Euclidean norm over vectors along dimension :math:`1` for normalization.</span>

<span class="sd">    Args:</span>
<span class="sd">        input: input tensor of any shape</span>
<span class="sd">        p (float): the exponent value in the norm formulation. Default: 2</span>
<span class="sd">        dim (int): the dimension to reduce. Default: 1</span>
<span class="sd">        eps (float): small value to avoid division by zero. Default: 1e-12</span>
<span class="sd">        out (Tensor, optional): the output tensor. If :attr:`out` is used, this</span>
<span class="sd">                                operation won&#39;t be differentiable.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_variadic</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">out</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">normalize</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">out</span><span class="p">),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">p</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="n">dim</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="n">eps</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">out</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">out</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">denom</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">norm</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span><span class="o">.</span><span class="n">clamp_min</span><span class="p">(</span><span class="n">eps</span><span class="p">)</span><span class="o">.</span><span class="n">expand_as</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="k">return</span> <span class="nb">input</span> <span class="o">/</span> <span class="n">denom</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">denom</span> <span class="o">=</span> <span class="nb">input</span><span class="o">.</span><span class="n">norm</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span><span class="o">.</span><span class="n">clamp_min_</span><span class="p">(</span><span class="n">eps</span><span class="p">)</span><span class="o">.</span><span class="n">expand_as</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">div</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">denom</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="n">out</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">assert_int_or_pair</span><span class="p">(</span><span class="n">arg</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">arg_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">message</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
    <span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="nb">int</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">arg</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">,</span> <span class="n">message</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">arg_name</span><span class="p">)</span>


<div class="viewcode-block" id="unfold"><a class="viewcode-back" href="../../../generated/torch.nn.functional.unfold.html#torch.nn.functional.unfold">[docs]</a><span class="k">def</span> <span class="nf">unfold</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Extracts sliding local blocks from a batched input tensor.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        Currently, only 4-D input tensors (batched image-like tensors) are</span>
<span class="sd">        supported.</span>

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

<span class="sd">        More than one element of the unfolded tensor may refer to a single</span>
<span class="sd">        memory location. As a result, in-place operations (especially ones that</span>
<span class="sd">        are vectorized) may result in incorrect behavior. If you need to write</span>
<span class="sd">        to the tensor, please clone it first.</span>


<span class="sd">    See :class:`torch.nn.Unfold` for details</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">unfold</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span> <span class="n">stride</span><span class="o">=</span><span class="n">stride</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">im2col</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">_pair</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">dilation</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">padding</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">stride</span><span class="p">))</span></div>


<div class="viewcode-block" id="fold"><a class="viewcode-back" href="../../../generated/torch.nn.functional.fold.html#torch.nn.functional.fold">[docs]</a><span class="k">def</span> <span class="nf">fold</span><span class="p">(</span>
    <span class="nb">input</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">output_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">kernel_size</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span>
    <span class="n">dilation</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
    <span class="n">padding</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
    <span class="n">stride</span><span class="p">:</span> <span class="n">BroadcastingList2</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;Combines an array of sliding local blocks into a large containing</span>
<span class="sd">    tensor.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        Currently, only unbatched (3D) or batched (4D) image-like output tensors are supported.</span>

<span class="sd">    See :class:`torch.nn.Fold` for details</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="nb">input</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">fold</span><span class="p">,</span> <span class="p">(</span><span class="nb">input</span><span class="p">,),</span> <span class="nb">input</span><span class="p">,</span> <span class="n">output_size</span><span class="p">,</span> <span class="n">kernel_size</span><span class="p">,</span> <span class="n">dilation</span><span class="o">=</span><span class="n">dilation</span><span class="p">,</span> <span class="n">padding</span><span class="o">=</span><span class="n">padding</span><span class="p">,</span> <span class="n">stride</span><span class="o">=</span><span class="n">stride</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_nn</span><span class="o">.</span><span class="n">col2im</span><span class="p">(</span>
        <span class="nb">input</span><span class="p">,</span> <span class="n">_pair</span><span class="p">(</span><span class="n">output_size</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">kernel_size</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">dilation</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">padding</span><span class="p">),</span> <span class="n">_pair</span><span class="p">(</span><span class="n">stride</span><span class="p">)</span>
    <span class="p">)</span></div>

<span class="c1">#</span>
<span class="c1"># multihead attention</span>
<span class="c1">#</span>

<span class="k">def</span> <span class="nf">_in_projection_packed</span><span class="p">(</span>
    <span class="n">q</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">k</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">v</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">w</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">b</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Performs the in-projection step of the attention operation, using packed weights.</span>
<span class="sd">    Output is a triple containing projection tensors for query, key and value.</span>

<span class="sd">    Args:</span>
<span class="sd">        q, k, v: query, key and value tensors to be projected. For self-attention,</span>
<span class="sd">            these are typically the same tensor; for encoder-decoder attention,</span>
<span class="sd">            k and v are typically the same tensor. (We take advantage of these</span>
<span class="sd">            identities for performance if they are present.) Regardless, q, k and v</span>
<span class="sd">            must share a common embedding dimension; otherwise their shapes may vary.</span>
<span class="sd">        w: projection weights for q, k and v, packed into a single tensor. Weights</span>
<span class="sd">            are packed along dimension 0, in q, k, v order.</span>
<span class="sd">        b: optional projection biases for q, k and v, packed into a single tensor</span>
<span class="sd">            in q, k, v order.</span>

<span class="sd">    Shape:</span>
<span class="sd">        Inputs:</span>
<span class="sd">        - q: :math:`(..., E)` where E is the embedding dimension</span>
<span class="sd">        - k: :math:`(..., E)` where E is the embedding dimension</span>
<span class="sd">        - v: :math:`(..., E)` where E is the embedding dimension</span>
<span class="sd">        - w: :math:`(E * 3, E)` where E is the embedding dimension</span>
<span class="sd">        - b: :math:`E * 3` where E is the embedding dimension</span>

<span class="sd">        Output:</span>
<span class="sd">        - in output list :math:`[q&#39;, k&#39;, v&#39;]`, each output tensor will have the</span>
<span class="sd">            same shape as the corresponding input tensor.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">E</span> <span class="o">=</span> <span class="n">q</span><span class="o">.</span><span class="n">size</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">k</span> <span class="ow">is</span> <span class="n">v</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">q</span> <span class="ow">is</span> <span class="n">k</span><span class="p">:</span>
            <span class="c1"># self-attention</span>
            <span class="n">proj</span> <span class="o">=</span> <span class="n">linear</span><span class="p">(</span><span class="n">q</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">b</span><span class="p">)</span>
            <span class="c1"># reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()</span>
            <span class="n">proj</span> <span class="o">=</span> <span class="n">proj</span><span class="o">.</span><span class="n">unflatten</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="n">E</span><span class="p">))</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span>
            <span class="k">return</span> <span class="n">proj</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">proj</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">proj</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># encoder-decoder attention</span>
            <span class="n">w_q</span><span class="p">,</span> <span class="n">w_kv</span> <span class="o">=</span> <span class="n">w</span><span class="o">.</span><span class="n">split</span><span class="p">([</span><span class="n">E</span><span class="p">,</span> <span class="n">E</span> <span class="o">*</span> <span class="mi">2</span><span class="p">])</span>
            <span class="k">if</span> <span class="n">b</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">b_q</span> <span class="o">=</span> <span class="n">b_kv</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">b_q</span><span class="p">,</span> <span class="n">b_kv</span> <span class="o">=</span> <span class="n">b</span><span class="o">.</span><span class="n">split</span><span class="p">([</span><span class="n">E</span><span class="p">,</span> <span class="n">E</span> <span class="o">*</span> <span class="mi">2</span><span class="p">])</span>
            <span class="n">q_proj</span> <span class="o">=</span> <span class="n">linear</span><span class="p">(</span><span class="n">q</span><span class="p">,</span> <span class="n">w_q</span><span class="p">,</span> <span class="n">b_q</span><span class="p">)</span>
            <span class="n">kv_proj</span> <span class="o">=</span> <span class="n">linear</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">w_kv</span><span class="p">,</span> <span class="n">b_kv</span><span class="p">)</span>
            <span class="c1"># reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()</span>
            <span class="n">kv_proj</span> <span class="o">=</span> <span class="n">kv_proj</span><span class="o">.</span><span class="n">unflatten</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="n">E</span><span class="p">))</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">)</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span>
            <span class="k">return</span> <span class="p">(</span><span class="n">q_proj</span><span class="p">,</span> <span class="n">kv_proj</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">kv_proj</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">w_q</span><span class="p">,</span> <span class="n">w_k</span><span class="p">,</span> <span class="n">w_v</span> <span class="o">=</span> <span class="n">w</span><span class="o">.</span><span class="n">chunk</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">b</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">b_q</span> <span class="o">=</span> <span class="n">b_k</span> <span class="o">=</span> <span class="n">b_v</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">b_q</span><span class="p">,</span> <span class="n">b_k</span><span class="p">,</span> <span class="n">b_v</span> <span class="o">=</span> <span class="n">b</span><span class="o">.</span><span class="n">chunk</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">linear</span><span class="p">(</span><span class="n">q</span><span class="p">,</span> <span class="n">w_q</span><span class="p">,</span> <span class="n">b_q</span><span class="p">),</span> <span class="n">linear</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">w_k</span><span class="p">,</span> <span class="n">b_k</span><span class="p">),</span> <span class="n">linear</span><span class="p">(</span><span class="n">v</span><span class="p">,</span> <span class="n">w_v</span><span class="p">,</span> <span class="n">b_v</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_in_projection</span><span class="p">(</span>
    <span class="n">q</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">k</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">v</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">w_q</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">w_k</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">w_v</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">b_q</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">b_k</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">b_v</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
<span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">Tensor</span><span class="p">]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Performs the in-projection step of the attention operation. This is simply</span>
<span class="sd">    a triple of linear projections, with shape constraints on the weights which</span>
<span class="sd">    ensure embedding dimension uniformity in the projected outputs.</span>
<span class="sd">    Output is a triple containing projection tensors for query, key and value.</span>

<span class="sd">    Args:</span>
<span class="sd">        q, k, v: query, key and value tensors to be projected.</span>
<span class="sd">        w_q, w_k, w_v: weights for q, k and v, respectively.</span>
<span class="sd">        b_q, b_k, b_v: optional biases for q, k and v, respectively.</span>

<span class="sd">    Shape:</span>
<span class="sd">        Inputs:</span>
<span class="sd">        - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any</span>
<span class="sd">            number of leading dimensions.</span>
<span class="sd">        - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any</span>
<span class="sd">            number of leading dimensions.</span>
<span class="sd">        - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any</span>
<span class="sd">            number of leading dimensions.</span>
<span class="sd">        - w_q: :math:`(Eq, Eq)`</span>
<span class="sd">        - w_k: :math:`(Eq, Ek)`</span>
<span class="sd">        - w_v: :math:`(Eq, Ev)`</span>
<span class="sd">        - b_q: :math:`(Eq)`</span>
<span class="sd">        - b_k: :math:`(Eq)`</span>
<span class="sd">        - b_v: :math:`(Eq)`</span>

<span class="sd">        Output: in output triple :math:`(q&#39;, k&#39;, v&#39;)`,</span>
<span class="sd">         - q&#39;: :math:`[Qdims..., Eq]`</span>
<span class="sd">         - k&#39;: :math:`[Kdims..., Eq]`</span>
<span class="sd">         - v&#39;: :math:`[Vdims..., Eq]`</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">Eq</span><span class="p">,</span> <span class="n">Ek</span><span class="p">,</span> <span class="n">Ev</span> <span class="o">=</span> <span class="n">q</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">k</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">v</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
    <span class="k">assert</span> <span class="n">w_q</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">Eq</span><span class="p">,</span> <span class="n">Eq</span><span class="p">),</span> <span class="sa">f</span><span class="s2">&quot;expecting query weights shape of </span><span class="si">{</span><span class="p">(</span><span class="n">Eq</span><span class="p">,</span> <span class="n">Eq</span><span class="p">)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">w_q</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">assert</span> <span class="n">w_k</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">Eq</span><span class="p">,</span> <span class="n">Ek</span><span class="p">),</span> <span class="sa">f</span><span class="s2">&quot;expecting key weights shape of </span><span class="si">{</span><span class="p">(</span><span class="n">Eq</span><span class="p">,</span> <span class="n">Ek</span><span class="p">)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">w_k</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">assert</span> <span class="n">w_v</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">Eq</span><span class="p">,</span> <span class="n">Ev</span><span class="p">),</span> <span class="sa">f</span><span class="s2">&quot;expecting value weights shape of </span><span class="si">{</span><span class="p">(</span><span class="n">Eq</span><span class="p">,</span> <span class="n">Ev</span><span class="p">)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">w_v</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">assert</span> <span class="n">b_q</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">b_q</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">Eq</span><span class="p">,),</span> <span class="sa">f</span><span class="s2">&quot;expecting query bias shape of </span><span class="si">{</span><span class="p">(</span><span class="n">Eq</span><span class="p">,)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">b_q</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">assert</span> <span class="n">b_k</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">b_k</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">Eq</span><span class="p">,),</span> <span class="sa">f</span><span class="s2">&quot;expecting key bias shape of </span><span class="si">{</span><span class="p">(</span><span class="n">Eq</span><span class="p">,)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">b_k</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">assert</span> <span class="n">b_v</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">b_v</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">Eq</span><span class="p">,),</span> <span class="sa">f</span><span class="s2">&quot;expecting value bias shape of </span><span class="si">{</span><span class="p">(</span><span class="n">Eq</span><span class="p">,)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">b_v</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">return</span> <span class="n">linear</span><span class="p">(</span><span class="n">q</span><span class="p">,</span> <span class="n">w_q</span><span class="p">,</span> <span class="n">b_q</span><span class="p">),</span> <span class="n">linear</span><span class="p">(</span><span class="n">k</span><span class="p">,</span> <span class="n">w_k</span><span class="p">,</span> <span class="n">b_k</span><span class="p">),</span> <span class="n">linear</span><span class="p">(</span><span class="n">v</span><span class="p">,</span> <span class="n">w_v</span><span class="p">,</span> <span class="n">b_v</span><span class="p">)</span>

<span class="n">scaled_dot_product_attention</span> <span class="o">=</span> <span class="n">_add_docstr</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">_nn</span><span class="o">.</span><span class="n">scaled_dot_product_attention</span><span class="p">,</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>
<span class="s2">scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False) -&gt; Tensor:</span>

<span class="s2">Computes scaled dot product attention on query, key and value tensors, using</span>
<span class="s2">an optional attention mask if passed, and applying dropout if a probability</span>
<span class="s2">greater than 0.0 is specified.</span>

<span class="s2">.. code-block:: python</span>

<span class="s2">    # Efficient implementation equivalent to the following:</span>
<span class="s2">    attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask</span>
<span class="s2">    attn_mask = attn_mask.masked_fill(not attn_mask, -float(&#39;inf&#39;)) if attn_mask.dtype==torch.bool else attn_mask</span>
<span class="s2">    attn_weight = torch.softmax((Q @ K.transpose(-2, -1) / math.sqrt(Q.size(-1))) + attn_mask, dim=-1)</span>
<span class="s2">    attn_weight = torch.dropout(attn_weight, dropout_p)</span>
<span class="s2">    return attn_weight @ V</span>

<span class="s2">.. warning:: This function is beta and subject to change.</span>

<span class="s2">Note:</span>

<span class="s2">    There are currently three supported implementations of scaled dot product attention:</span>

<span class="s2">        - `FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness`_</span>
<span class="s2">        - `Memory-Efficient Attention`_</span>
<span class="s2">        - A PyTorch implementation defined in C++ matching the above formulation</span>

<span class="s2">    The function may call optimized kernels for improved performance when using the CUDA backend.</span>
<span class="s2">    For all other backends, the PyTorch implementation will be used.</span>

<span class="s2">    All implementations are enabled by default. Scaled dot product attention attempts to automatically select the</span>
<span class="s2">    most optimal implementation based on the inputs. In order to provide more fine-grained control over what implementation</span>
<span class="s2">    is used, the following functions are provided for enabling and disabling implementations.</span>
<span class="s2">    The context manager is the preferred mechanism:</span>

<span class="s2">        - :func:`torch.backends.cuda.sdp_kernel`: A context manager used to enable/disable any of the implementations.</span>
<span class="s2">        - :func:`torch.backends.cuda.enable_flash_sdp`: Enables or Disables FlashAttention.</span>
<span class="s2">        - :func:`torch.backends.cuda.enable_mem_efficient_sdp`: Enables or Disables Memory-Efficient Attention.</span>
<span class="s2">        - :func:`torch.backends.cuda.enable_math_sdp`: Enables or Disables the PyTorch C++ implementation.</span>

<span class="s2">    Each of the fused kernels has specific input limitations. If the user requires the use of a specific fused implementation,</span>
<span class="s2">    disable the PyTorch C++ implementation using :func:`torch.backends.cuda.sdp_kernel`.</span>
<span class="s2">    In the event that a fused implementation is not available, an error will be raised with the</span>
<span class="s2">    reasons why the fused implementation cannot run.</span>

<span class="s2">    Due to the nature of fusing floating point operations, the output of this function may be different</span>
<span class="s2">    depending on what backend kernel is chosen.</span>
<span class="s2">    The c++ implementation supports torch.float64 and can be used when higher precision is required.</span>
<span class="s2">    For more information please see :doc:`/notes/numerical_accuracy`</span>

<span class="s2">Note:</span>
<span class="s2">    </span><span class="si">{cudnn_reproducibility_note}</span><span class="s2"></span>
<span class="s2">&quot;&quot;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="o">**</span><span class="n">reproducibility_notes</span><span class="p">)</span>
    <span class="o">+</span> <span class="sa">r</span><span class="s2">&quot;&quot;&quot;</span>

<span class="s2">Args:</span>
<span class="s2">    query (Tensor): Query tensor; shape :math:`(N, ..., L, E)`.</span>
<span class="s2">    key (Tensor): Key tensor; shape :math:`(N, ..., S, E)`.</span>
<span class="s2">    value (Tensor): Value tensor; shape :math:`(N, ..., S, Ev)`.</span>
<span class="s2">    attn_mask (optional Tensor): Attention mask; shape :math:`(N, ..., L, S)`. Two types of masks are supported.</span>
<span class="s2">        A boolean mask where a value of True indicates that the element *should* take part in attention.</span>
<span class="s2">        A float mask of the same type as query, key, value that is added to the attention score.</span>
<span class="s2">    dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied</span>
<span class="s2">    is_causal (bool): If true, assumes causal attention masking and errors if both attn_mask and is_causal</span>
<span class="s2">        are set.</span>


<span class="s2">Returns:</span>
<span class="s2">    output (Tensor): Attention output; shape :math:`(N, ..., L, Ev)`.</span>

<span class="s2">Shape legend:</span>
<span class="s2">    - :math:`N: \text{Batch size} ... : \text{Any number of other batch dimensions (optional)}`</span>
<span class="s2">    - :math:`S: \text{Source sequence length}`</span>
<span class="s2">    - :math:`L: \text{Target sequence length}`</span>
<span class="s2">    - :math:`E: \text{Embedding dimension of the query and key}`</span>
<span class="s2">    - :math:`Ev: \text{Embedding dimension of the value}`</span>

<span class="s2">Examples::</span>

<span class="s2">    &gt;&gt;&gt; # Optionally use the context manager to ensure one of the fused kerenels is run</span>
<span class="s2">    &gt;&gt;&gt; query = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=&quot;cuda&quot;)</span>
<span class="s2">    &gt;&gt;&gt; key = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=&quot;cuda&quot;)</span>
<span class="s2">    &gt;&gt;&gt; value = torch.rand(32, 8, 128, 64, dtype=torch.float16, device=&quot;cuda&quot;)</span>
<span class="s2">    &gt;&gt;&gt; with torch.backends.cuda.sdp_kernel(enable_math=False):</span>
<span class="s2">    &gt;&gt;&gt;     F.scaled_dot_product_attention(query,key,value)</span>

<span class="s2">.. _FlashAttention\: Fast and Memory-Efficient Exact Attention with IO-Awareness:</span>
<span class="s2">    https://arxiv.org/abs/2205.14135</span>
<span class="s2">.. _Memory-Efficient Attention:</span>
<span class="s2">    https://github.com/facebookresearch/xformers</span>

<span class="s2">&quot;&quot;&quot;</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">_mha_shape_check</span><span class="p">(</span><span class="n">query</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">key</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
                     <span class="n">key_padding_mask</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">attn_mask</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span> <span class="n">num_heads</span><span class="p">:</span> <span class="nb">int</span><span class="p">):</span>
    <span class="c1"># Verifies the expected shape for `query, `key`, `value`, `key_padding_mask` and `attn_mask`</span>
    <span class="c1"># and returns if the input is batched or not.</span>
    <span class="c1"># Raises an error if `query` is not 2-D (unbatched) or 3-D (batched) tensor.</span>

    <span class="c1"># Shape check.</span>
    <span class="k">if</span> <span class="n">query</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
        <span class="c1"># Batched Inputs</span>
        <span class="n">is_batched</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">assert</span> <span class="n">key</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">value</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span><span class="p">,</span> \
            <span class="p">(</span><span class="s2">&quot;For batched (3-D) `query`, expected `key` and `value` to be 3-D&quot;</span>
             <span class="sa">f</span><span class="s2">&quot; but found </span><span class="si">{</span><span class="n">key</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D and </span><span class="si">{</span><span class="n">value</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D tensors respectively&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">2</span><span class="p">,</span> \
                <span class="p">(</span><span class="s2">&quot;For batched (3-D) `query`, expected `key_padding_mask` to be `None` or 2-D&quot;</span>
                 <span class="sa">f</span><span class="s2">&quot; but found </span><span class="si">{</span><span class="n">key_padding_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D tensor instead&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> \
                <span class="p">(</span><span class="s2">&quot;For batched (3-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D&quot;</span>
                 <span class="sa">f</span><span class="s2">&quot; but found </span><span class="si">{</span><span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D tensor instead&quot;</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">query</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="c1"># Unbatched Inputs</span>
        <span class="n">is_batched</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="k">assert</span> <span class="n">key</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">2</span> <span class="ow">and</span> <span class="n">value</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">2</span><span class="p">,</span> \
            <span class="p">(</span><span class="s2">&quot;For unbatched (2-D) `query`, expected `key` and `value` to be 2-D&quot;</span>
             <span class="sa">f</span><span class="s2">&quot; but found </span><span class="si">{</span><span class="n">key</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D and </span><span class="si">{</span><span class="n">value</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D tensors respectively&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">1</span><span class="p">,</span> \
                <span class="p">(</span><span class="s2">&quot;For unbatched (2-D) `query`, expected `key_padding_mask` to be `None` or 1-D&quot;</span>
                 <span class="sa">f</span><span class="s2">&quot; but found </span><span class="si">{</span><span class="n">key_padding_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D tensor instead&quot;</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="ow">in</span> <span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> \
                <span class="p">(</span><span class="s2">&quot;For unbatched (2-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D&quot;</span>
                 <span class="sa">f</span><span class="s2">&quot; but found </span><span class="si">{</span><span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D tensor instead&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
                <span class="n">expected_shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">num_heads</span><span class="p">,</span> <span class="n">query</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">key</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
                <span class="k">assert</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="n">expected_shape</span><span class="p">,</span> \
                    <span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Expected `attn_mask` shape to be </span><span class="si">{</span><span class="n">expected_shape</span><span class="si">}</span><span class="s2"> but got </span><span class="si">{</span><span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</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="k">raise</span> <span class="ne">AssertionError</span><span class="p">(</span>
            <span class="sa">f</span><span class="s2">&quot;query should be unbatched 2D or batched 3D tensor but received </span><span class="si">{</span><span class="n">query</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2">-D query tensor&quot;</span><span class="p">)</span>

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

<span class="k">def</span> <span class="nf">_canonical_mask</span><span class="p">(</span>
        <span class="n">mask</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
        <span class="n">mask_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
        <span class="n">other_type</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">DType</span><span class="p">],</span>
        <span class="n">other_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
        <span class="n">target_type</span><span class="p">:</span> <span class="n">DType</span><span class="p">,</span>
        <span class="n">check_other</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="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]:</span>

    <span class="k">if</span> <span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">_mask_dtype</span> <span class="o">=</span> <span class="n">mask</span><span class="o">.</span><span class="n">dtype</span>
        <span class="n">_mask_is_float</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">(</span><span class="n">mask</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">_mask_dtype</span> <span class="o">!=</span> <span class="n">torch</span><span class="o">.</span><span class="n">bool</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">_mask_is_float</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">AssertionError</span><span class="p">(</span>
                <span class="sa">f</span><span class="s2">&quot;only bool and floating types of </span><span class="si">{</span><span class="n">mask_name</span><span class="si">}</span><span class="s2"> are supported&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">check_other</span> <span class="ow">and</span> <span class="n">other_type</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">_mask_dtype</span> <span class="o">!=</span> <span class="n">other_type</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;Support for mismatched </span><span class="si">{</span><span class="n">mask_name</span><span class="si">}</span><span class="s2"> and </span><span class="si">{</span><span class="n">other_name</span><span class="si">}</span><span class="s2"> &quot;</span>
                    <span class="s2">&quot;is deprecated. Use same type for both instead.&quot;</span>
                <span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">_mask_is_float</span><span class="p">:</span>
            <span class="n">mask</span> <span class="o">=</span> <span class="p">(</span>
                <span class="n">torch</span><span class="o">.</span><span class="n">zeros_like</span><span class="p">(</span><span class="n">mask</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">target_type</span><span class="p">)</span>
                <span class="o">.</span><span class="n">masked_fill_</span><span class="p">(</span><span class="n">mask</span><span class="p">,</span> <span class="nb">float</span><span class="p">(</span><span class="s2">&quot;-inf&quot;</span><span class="p">))</span>
            <span class="p">)</span>
    <span class="k">return</span> <span class="n">mask</span>

<span class="k">def</span> <span class="nf">_none_or_dtype</span><span class="p">(</span><span class="nb">input</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Optional</span><span class="p">[</span><span class="n">DType</span><span class="p">]:</span>
    <span class="k">if</span> <span class="nb">input</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">return</span> <span class="kc">None</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="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">input</span><span class="o">.</span><span class="n">dtype</span>
    <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;input to _none_or_dtype() must be None or torch.Tensor&quot;</span><span class="p">)</span>

<span class="k">def</span> <span class="nf">multi_head_attention_forward</span><span class="p">(</span>
    <span class="n">query</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">key</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">value</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">embed_dim_to_check</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
    <span class="n">num_heads</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
    <span class="n">in_proj_weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">in_proj_bias</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">bias_k</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">bias_v</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">add_zero_attn</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span>
    <span class="n">dropout_p</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
    <span class="n">out_proj_weight</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">,</span>
    <span class="n">out_proj_bias</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">],</span>
    <span class="n">training</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">key_padding_mask</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">need_weights</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">attn_mask</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">use_separate_proj_weight</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">q_proj_weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">k_proj_weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">v_proj_weight</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">static_k</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">static_v</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
    <span class="n">average_attn_weights</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">is_causal</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="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]]:</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Args:</span>
<span class="sd">        query, key, value: map a query and a set of key-value pairs to an output.</span>
<span class="sd">            See &quot;Attention Is All You Need&quot; for more details.</span>
<span class="sd">        embed_dim_to_check: total dimension of the model.</span>
<span class="sd">        num_heads: parallel attention heads.</span>
<span class="sd">        in_proj_weight, in_proj_bias: input projection weight and bias.</span>
<span class="sd">        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.</span>
<span class="sd">        add_zero_attn: add a new batch of zeros to the key and</span>
<span class="sd">                       value sequences at dim=1.</span>
<span class="sd">        dropout_p: probability of an element to be zeroed.</span>
<span class="sd">        out_proj_weight, out_proj_bias: the output projection weight and bias.</span>
<span class="sd">        training: apply dropout if is ``True``.</span>
<span class="sd">        key_padding_mask: if provided, specified padding elements in the key will</span>
<span class="sd">            be ignored by the attention. This is an binary mask. When the value is True,</span>
<span class="sd">            the corresponding value on the attention layer will be filled with -inf.</span>
<span class="sd">        need_weights: output attn_output_weights.</span>
<span class="sd">            Default: `True`</span>
<span class="sd">            Note: `needs_weight` defaults to `True`, but should be set to `False`</span>
<span class="sd">            For best performance when attention weights are not nedeeded.</span>
<span class="sd">            *Setting needs_weights to `True`</span>
<span class="sd">            leads to a significant performance degradation.*</span>
<span class="sd">        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all</span>
<span class="sd">            the batches while a 3D mask allows to specify a different mask for the entries of each batch.</span>
<span class="sd">        is_causal: If specified, applies a causal mask as attention mask, and ignores</span>
<span class="sd">            attn_mask for computing scaled dot product attention.</span>
<span class="sd">            Default: ``False``.</span>
<span class="sd">            .. warning::</span>
<span class="sd">                is_causal is provides a hint that the attn_mask is the</span>
<span class="sd">                causal mask.Providing incorrect hints can result in</span>
<span class="sd">                incorrect execution, including forward and backward</span>
<span class="sd">                compatibility.</span>
<span class="sd">        use_separate_proj_weight: the function accept the proj. weights for query, key,</span>
<span class="sd">            and value in different forms. If false, in_proj_weight will be used, which is</span>
<span class="sd">            a combination of q_proj_weight, k_proj_weight, v_proj_weight.</span>
<span class="sd">        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.</span>
<span class="sd">        static_k, static_v: static key and value used for attention operators.</span>
<span class="sd">        average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across heads.</span>
<span class="sd">            Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an effect</span>
<span class="sd">            when ``need_weights=True.``. Default: True</span>


<span class="sd">    Shape:</span>
<span class="sd">        Inputs:</span>
<span class="sd">        - query: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is</span>
<span class="sd">          the embedding dimension.</span>
<span class="sd">        - key: :math:`(S, E)` or :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is</span>
<span class="sd">          the embedding dimension.</span>
<span class="sd">        - value: :math:`(S, E)` or :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is</span>
<span class="sd">          the embedding dimension.</span>
<span class="sd">        - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.</span>
<span class="sd">          If a FloatTensor is provided, it will be directly added to the value.</span>
<span class="sd">          If a BoolTensor is provided, the positions with the</span>
<span class="sd">          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.</span>
<span class="sd">        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.</span>
<span class="sd">          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,</span>
<span class="sd">          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked</span>
<span class="sd">          positions. If a BoolTensor is provided, positions with ``True``</span>
<span class="sd">          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor</span>
<span class="sd">          is provided, it will be added to the attention weight.</span>
<span class="sd">        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,</span>
<span class="sd">          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.</span>
<span class="sd">        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,</span>
<span class="sd">          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.</span>

<span class="sd">        Outputs:</span>
<span class="sd">        - attn_output: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size,</span>
<span class="sd">          E is the embedding dimension.</span>
<span class="sd">        - attn_output_weights: Only returned when ``need_weights=True``. If ``average_attn_weights=True``, returns</span>
<span class="sd">          attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or</span>
<span class="sd">          :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and</span>
<span class="sd">          :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per</span>
<span class="sd">          head of shape :math:`(num_heads, L, S)` when input is unbatched or :math:`(N, num_heads, L, S)`.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">tens_ops</span> <span class="o">=</span> <span class="p">(</span><span class="n">query</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">in_proj_weight</span><span class="p">,</span> <span class="n">in_proj_bias</span><span class="p">,</span> <span class="n">bias_k</span><span class="p">,</span> <span class="n">bias_v</span><span class="p">,</span> <span class="n">out_proj_weight</span><span class="p">,</span> <span class="n">out_proj_bias</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tens_ops</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">multi_head_attention_forward</span><span class="p">,</span>
            <span class="n">tens_ops</span><span class="p">,</span>
            <span class="n">query</span><span class="p">,</span>
            <span class="n">key</span><span class="p">,</span>
            <span class="n">value</span><span class="p">,</span>
            <span class="n">embed_dim_to_check</span><span class="p">,</span>
            <span class="n">num_heads</span><span class="p">,</span>
            <span class="n">in_proj_weight</span><span class="p">,</span>
            <span class="n">in_proj_bias</span><span class="p">,</span>
            <span class="n">bias_k</span><span class="p">,</span>
            <span class="n">bias_v</span><span class="p">,</span>
            <span class="n">add_zero_attn</span><span class="p">,</span>
            <span class="n">dropout_p</span><span class="p">,</span>
            <span class="n">out_proj_weight</span><span class="p">,</span>
            <span class="n">out_proj_bias</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">key_padding_mask</span><span class="o">=</span><span class="n">key_padding_mask</span><span class="p">,</span>
            <span class="n">need_weights</span><span class="o">=</span><span class="n">need_weights</span><span class="p">,</span>
            <span class="n">attn_mask</span><span class="o">=</span><span class="n">attn_mask</span><span class="p">,</span>
            <span class="n">is_causal</span><span class="o">=</span><span class="n">is_causal</span><span class="p">,</span>
            <span class="n">use_separate_proj_weight</span><span class="o">=</span><span class="n">use_separate_proj_weight</span><span class="p">,</span>
            <span class="n">q_proj_weight</span><span class="o">=</span><span class="n">q_proj_weight</span><span class="p">,</span>
            <span class="n">k_proj_weight</span><span class="o">=</span><span class="n">k_proj_weight</span><span class="p">,</span>
            <span class="n">v_proj_weight</span><span class="o">=</span><span class="n">v_proj_weight</span><span class="p">,</span>
            <span class="n">static_k</span><span class="o">=</span><span class="n">static_k</span><span class="p">,</span>
            <span class="n">static_v</span><span class="o">=</span><span class="n">static_v</span><span class="p">,</span>
            <span class="n">average_attn_weights</span><span class="o">=</span><span class="n">average_attn_weights</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="n">is_batched</span> <span class="o">=</span> <span class="n">_mha_shape_check</span><span class="p">(</span><span class="n">query</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">key_padding_mask</span><span class="p">,</span> <span class="n">attn_mask</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">)</span>

    <span class="c1"># For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input</span>
    <span class="c1"># is batched, run the computation and before returning squeeze the</span>
    <span class="c1"># batch dimension so that the output doesn&#39;t carry this temporary batch dimension.</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
        <span class="c1"># unsqueeze if the input is unbatched</span>
        <span class="n">query</span> <span class="o">=</span> <span class="n">query</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">key</span> <span class="o">=</span> <span class="n">key</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">value</span> <span class="o">=</span> <span class="n">value</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>

    <span class="c1"># set up shape vars</span>
    <span class="n">tgt_len</span><span class="p">,</span> <span class="n">bsz</span><span class="p">,</span> <span class="n">embed_dim</span> <span class="o">=</span> <span class="n">query</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">src_len</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">key</span><span class="o">.</span><span class="n">shape</span>

    <span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">_canonical_mask</span><span class="p">(</span>
        <span class="n">mask</span><span class="o">=</span><span class="n">key_padding_mask</span><span class="p">,</span>
        <span class="n">mask_name</span><span class="o">=</span><span class="s2">&quot;key_padding_mask&quot;</span><span class="p">,</span>
        <span class="n">other_type</span><span class="o">=</span><span class="n">_none_or_dtype</span><span class="p">(</span><span class="n">attn_mask</span><span class="p">),</span>
        <span class="n">other_name</span><span class="o">=</span><span class="s2">&quot;attn_mask&quot;</span><span class="p">,</span>
        <span class="n">target_type</span><span class="o">=</span><span class="n">query</span><span class="o">.</span><span class="n">dtype</span>
    <span class="p">)</span>

    <span class="k">if</span> <span class="n">is_causal</span> <span class="ow">and</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span>
            <span class="s2">&quot;Need attn_mask if specifying the is_causal hint. &quot;</span>
            <span class="s2">&quot;You may use the Transformer module method &quot;</span>
            <span class="s2">&quot;`generate_square_subsequent_mask` to create this mask.&quot;</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">is_causal</span> <span class="ow">and</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">need_weights</span><span class="p">:</span>
        <span class="c1"># when we have a kpm or need weights, we need attn_mask</span>
        <span class="c1"># Otherwise, we use the is_causal hint go as is_causal</span>
        <span class="c1"># indicator to SDPA.</span>
        <span class="n">attn_mask</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">_canonical_mask</span><span class="p">(</span>
            <span class="n">mask</span><span class="o">=</span><span class="n">attn_mask</span><span class="p">,</span>
            <span class="n">mask_name</span><span class="o">=</span><span class="s2">&quot;attn_mask&quot;</span><span class="p">,</span>
            <span class="n">other_type</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
            <span class="n">other_name</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
            <span class="n">target_type</span><span class="o">=</span><span class="n">query</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span>
            <span class="n">check_other</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="p">)</span>


        <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="c1"># We have the attn_mask, and use that to merge kpm into it.</span>
            <span class="c1"># Turn off use of is_causal hint, as the merged mask is no</span>
            <span class="c1"># longer causal.</span>
            <span class="n">is_causal</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">assert</span> <span class="n">embed_dim</span> <span class="o">==</span> <span class="n">embed_dim_to_check</span><span class="p">,</span> \
        <span class="sa">f</span><span class="s2">&quot;was expecting embedding dimension of </span><span class="si">{</span><span class="n">embed_dim_to_check</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">embed_dim</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">embed_dim</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="c1"># embed_dim can be a tensor when JIT tracing</span>
        <span class="n">head_dim</span> <span class="o">=</span> <span class="n">embed_dim</span><span class="o">.</span><span class="n">div</span><span class="p">(</span><span class="n">num_heads</span><span class="p">,</span> <span class="n">rounding_mode</span><span class="o">=</span><span class="s1">&#39;trunc&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">head_dim</span> <span class="o">=</span> <span class="n">embed_dim</span> <span class="o">//</span> <span class="n">num_heads</span>
    <span class="k">assert</span> <span class="n">head_dim</span> <span class="o">*</span> <span class="n">num_heads</span> <span class="o">==</span> <span class="n">embed_dim</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;embed_dim </span><span class="si">{</span><span class="n">embed_dim</span><span class="si">}</span><span class="s2"> not divisible by num_heads </span><span class="si">{</span><span class="n">num_heads</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">if</span> <span class="n">use_separate_proj_weight</span><span class="p">:</span>
        <span class="c1"># allow MHA to have different embedding dimensions when separate projection weights are used</span>
        <span class="k">assert</span> <span class="n">key</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span> <span class="o">==</span> <span class="n">value</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">],</span> \
            <span class="sa">f</span><span class="s2">&quot;key&#39;s sequence and batch dims </span><span class="si">{</span><span class="n">key</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span><span class="si">}</span><span class="s2"> do not match value&#39;s </span><span class="si">{</span><span class="n">value</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">]</span><span class="si">}</span><span class="s2">&quot;</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">key</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="n">value</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="sa">f</span><span class="s2">&quot;key shape </span><span class="si">{</span><span class="n">key</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2"> does not match value shape </span><span class="si">{</span><span class="n">value</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>

    <span class="c1">#</span>
    <span class="c1"># compute in-projection</span>
    <span class="c1">#</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">use_separate_proj_weight</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">in_proj_weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;use_separate_proj_weight is False but in_proj_weight is None&quot;</span>
        <span class="n">q</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="o">=</span> <span class="n">_in_projection_packed</span><span class="p">(</span><span class="n">query</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">in_proj_weight</span><span class="p">,</span> <span class="n">in_proj_bias</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">q_proj_weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;use_separate_proj_weight is True but q_proj_weight is None&quot;</span>
        <span class="k">assert</span> <span class="n">k_proj_weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;use_separate_proj_weight is True but k_proj_weight is None&quot;</span>
        <span class="k">assert</span> <span class="n">v_proj_weight</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;use_separate_proj_weight is True but v_proj_weight is None&quot;</span>
        <span class="k">if</span> <span class="n">in_proj_bias</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">b_q</span> <span class="o">=</span> <span class="n">b_k</span> <span class="o">=</span> <span class="n">b_v</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">b_q</span><span class="p">,</span> <span class="n">b_k</span><span class="p">,</span> <span class="n">b_v</span> <span class="o">=</span> <span class="n">in_proj_bias</span><span class="o">.</span><span class="n">chunk</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
        <span class="n">q</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="o">=</span> <span class="n">_in_projection</span><span class="p">(</span><span class="n">query</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">q_proj_weight</span><span class="p">,</span> <span class="n">k_proj_weight</span><span class="p">,</span> <span class="n">v_proj_weight</span><span class="p">,</span> <span class="n">b_q</span><span class="p">,</span> <span class="n">b_k</span><span class="p">,</span> <span class="n">b_v</span><span class="p">)</span>

    <span class="c1"># prep attention mask</span>

    <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">_canonical_mask</span><span class="p">(</span>
        <span class="n">mask</span><span class="o">=</span><span class="n">attn_mask</span><span class="p">,</span>
        <span class="n">mask_name</span><span class="o">=</span><span class="s2">&quot;attn_mask&quot;</span><span class="p">,</span>
        <span class="n">other_type</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">other_name</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
        <span class="n">target_type</span><span class="o">=</span><span class="n">q</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span>
        <span class="n">check_other</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
    <span class="p">)</span>

    <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># ensure attn_mask&#39;s dim is 3</span>
        <span class="k">if</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
            <span class="n">correct_2d_size</span> <span class="o">=</span> <span class="p">(</span><span class="n">tgt_len</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</span> <span class="o">!=</span> <span class="n">correct_2d_size</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The shape of the 2D attn_mask is </span><span class="si">{</span><span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">, but should be </span><span class="si">{</span><span class="n">correct_2d_size</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
            <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
            <span class="n">correct_3d_size</span> <span class="o">=</span> <span class="p">(</span><span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">tgt_len</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</span> <span class="o">!=</span> <span class="n">correct_3d_size</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;The shape of the 3D attn_mask is </span><span class="si">{</span><span class="n">attn_mask</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">, but should be </span><span class="si">{</span><span class="n">correct_3d_size</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="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;attn_mask&#39;s dimension </span><span class="si">{</span><span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span><span class="si">}</span><span class="s2"> is not supported&quot;</span><span class="p">)</span>

    <span class="c1"># add bias along batch dimension (currently second)</span>
    <span class="k">if</span> <span class="n">bias_k</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">bias_v</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">static_k</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;bias cannot be added to static key.&quot;</span>
        <span class="k">assert</span> <span class="n">static_v</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;bias cannot be added to static value.&quot;</span>
        <span class="n">k</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">k</span><span class="p">,</span> <span class="n">bias_k</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">bsz</span><span class="p">,</span> <span class="mi">1</span><span class="p">)])</span>
        <span class="n">v</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">v</span><span class="p">,</span> <span class="n">bias_v</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">bsz</span><span class="p">,</span> <span class="mi">1</span><span class="p">)])</span>
        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">pad</span><span class="p">(</span><span class="n">attn_mask</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">pad</span><span class="p">(</span><span class="n">key_padding_mask</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">bias_k</span> <span class="ow">is</span> <span class="kc">None</span>
        <span class="k">assert</span> <span class="n">bias_v</span> <span class="ow">is</span> <span class="kc">None</span>

    <span class="c1">#</span>
    <span class="c1"># reshape q, k, v for multihead attention and make em batch first</span>
    <span class="c1">#</span>
    <span class="n">q</span> <span class="o">=</span> <span class="n">q</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">tgt_len</span><span class="p">,</span> <span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">static_k</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">k</span> <span class="o">=</span> <span class="n">k</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">k</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># TODO finish disentangling control flow so we don&#39;t do in-projections when statics are passed</span>
        <span class="k">assert</span> <span class="n">static_k</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">==</span> <span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> \
            <span class="sa">f</span><span class="s2">&quot;expecting static_k.size(0) of </span><span class="si">{</span><span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">static_k</span><span class="o">.</span><span class="n">size</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="k">assert</span> <span class="n">static_k</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span> <span class="o">==</span> <span class="n">head_dim</span><span class="p">,</span> \
            <span class="sa">f</span><span class="s2">&quot;expecting static_k.size(2) of </span><span class="si">{</span><span class="n">head_dim</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">static_k</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span>
        <span class="n">k</span> <span class="o">=</span> <span class="n">static_k</span>
    <span class="k">if</span> <span class="n">static_v</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">v</span> <span class="o">=</span> <span class="n">v</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">v</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># TODO finish disentangling control flow so we don&#39;t do in-projections when statics are passed</span>
        <span class="k">assert</span> <span class="n">static_v</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">==</span> <span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> \
            <span class="sa">f</span><span class="s2">&quot;expecting static_v.size(0) of </span><span class="si">{</span><span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">static_v</span><span class="o">.</span><span class="n">size</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="k">assert</span> <span class="n">static_v</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span> <span class="o">==</span> <span class="n">head_dim</span><span class="p">,</span> \
            <span class="sa">f</span><span class="s2">&quot;expecting static_v.size(2) of </span><span class="si">{</span><span class="n">head_dim</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">static_v</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">2</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">static_v</span>

    <span class="c1"># add zero attention along batch dimension (now first)</span>
    <span class="k">if</span> <span class="n">add_zero_attn</span><span class="p">:</span>
        <span class="n">zero_attn_shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span>
        <span class="n">k</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">k</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">zero_attn_shape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">k</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">k</span><span class="o">.</span><span class="n">device</span><span class="p">)],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
        <span class="n">v</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">([</span><span class="n">v</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">zero_attn_shape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">v</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">v</span><span class="o">.</span><span class="n">device</span><span class="p">)],</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">pad</span><span class="p">(</span><span class="n">attn_mask</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span>
        <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">pad</span><span class="p">(</span><span class="n">key_padding_mask</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span>

    <span class="c1"># update source sequence length after adjustments</span>
    <span class="n">src_len</span> <span class="o">=</span> <span class="n">k</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>

    <span class="c1"># merge key padding and attention masks</span>
    <span class="k">if</span> <span class="n">key_padding_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">src_len</span><span class="p">),</span> \
            <span class="sa">f</span><span class="s2">&quot;expecting key_padding_mask shape of </span><span class="si">{</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span><span class="si">}</span><span class="s2">, but got </span><span class="si">{</span><span class="n">key_padding_mask</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
        <span class="n">key_padding_mask</span> <span class="o">=</span> <span class="n">key_padding_mask</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span><span class="o">.</span>   \
            <span class="n">expand</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">bsz</span> <span class="o">*</span> <span class="n">num_heads</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">key_padding_mask</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">attn_mask</span> <span class="o">+</span> <span class="n">key_padding_mask</span>

    <span class="c1"># adjust dropout probability</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">training</span><span class="p">:</span>
        <span class="n">dropout_p</span> <span class="o">=</span> <span class="mf">0.0</span>

    <span class="c1">#</span>
    <span class="c1"># (deep breath) calculate attention and out projection</span>
    <span class="c1">#</span>

    <span class="k">if</span> <span class="n">need_weights</span><span class="p">:</span>
        <span class="n">B</span><span class="p">,</span> <span class="n">Nt</span><span class="p">,</span> <span class="n">E</span> <span class="o">=</span> <span class="n">q</span><span class="o">.</span><span class="n">shape</span>
        <span class="n">q_scaled</span> <span class="o">=</span> <span class="n">q</span> <span class="o">/</span> <span class="n">math</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">E</span><span class="p">)</span>

        <span class="k">assert</span> <span class="ow">not</span> <span class="p">(</span><span class="n">is_causal</span> <span class="ow">and</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">),</span> <span class="s2">&quot;FIXME: is_causal not implemented for need_weights&quot;</span>

        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">attn_output_weights</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">baddbmm</span><span class="p">(</span><span class="n">attn_mask</span><span class="p">,</span> <span class="n">q_scaled</span><span class="p">,</span> <span class="n">k</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="o">-</span><span class="mi">2</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">attn_output_weights</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">bmm</span><span class="p">(</span><span class="n">q_scaled</span><span class="p">,</span> <span class="n">k</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">))</span>
        <span class="n">attn_output_weights</span> <span class="o">=</span> <span class="n">softmax</span><span class="p">(</span><span class="n">attn_output_weights</span><span class="p">,</span> <span class="n">dim</span><span class="o">=-</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">dropout_p</span> <span class="o">&gt;</span> <span class="mf">0.0</span><span class="p">:</span>
            <span class="n">attn_output_weights</span> <span class="o">=</span> <span class="n">dropout</span><span class="p">(</span><span class="n">attn_output_weights</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="n">dropout_p</span><span class="p">)</span>

        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">bmm</span><span class="p">(</span><span class="n">attn_output_weights</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span>

        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">transpose</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">tgt_len</span> <span class="o">*</span> <span class="n">bsz</span><span class="p">,</span> <span class="n">embed_dim</span><span class="p">)</span>
        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">linear</span><span class="p">(</span><span class="n">attn_output</span><span class="p">,</span> <span class="n">out_proj_weight</span><span class="p">,</span> <span class="n">out_proj_bias</span><span class="p">)</span>
        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">tgt_len</span><span class="p">,</span> <span class="n">bsz</span><span class="p">,</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span>

        <span class="c1"># optionally average attention weights over heads</span>
        <span class="n">attn_output_weights</span> <span class="o">=</span> <span class="n">attn_output_weights</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">tgt_len</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">average_attn_weights</span><span class="p">:</span>
            <span class="n">attn_output_weights</span> <span class="o">=</span> <span class="n">attn_output_weights</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
            <span class="c1"># squeeze the output if input was unbatched</span>
            <span class="n">attn_output</span> <span class="o">=</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
            <span class="n">attn_output_weights</span> <span class="o">=</span> <span class="n">attn_output_weights</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">attn_output</span><span class="p">,</span> <span class="n">attn_output_weights</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># attn_mask can be either (L,S) or (N*num_heads, L, S)</span>
        <span class="c1"># if attn_mask&#39;s shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)</span>
        <span class="c1"># in order to match the input for SDPA of (N, num_heads, L, S)</span>
        <span class="k">if</span> <span class="n">attn_mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">dim</span><span class="p">()</span> <span class="o">==</span> <span class="mi">3</span><span class="p">:</span>
                <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">unsqueeze</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">attn_mask</span> <span class="o">=</span> <span class="n">attn_mask</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="n">src_len</span><span class="p">)</span>

        <span class="n">q</span> <span class="o">=</span> <span class="n">q</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">tgt_len</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span>
        <span class="n">k</span> <span class="o">=</span> <span class="n">k</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">src_len</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span>
        <span class="n">v</span> <span class="o">=</span> <span class="n">v</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span><span class="p">,</span> <span class="n">num_heads</span><span class="p">,</span> <span class="n">src_len</span><span class="p">,</span> <span class="n">head_dim</span><span class="p">)</span>

        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">scaled_dot_product_attention</span><span class="p">(</span><span class="n">q</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">attn_mask</span><span class="p">,</span> <span class="n">dropout_p</span><span class="p">,</span> <span class="n">is_causal</span><span class="p">)</span>
        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">permute</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span><span class="o">.</span><span class="n">contiguous</span><span class="p">()</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">bsz</span> <span class="o">*</span> <span class="n">tgt_len</span><span class="p">,</span> <span class="n">embed_dim</span><span class="p">)</span>

        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">linear</span><span class="p">(</span><span class="n">attn_output</span><span class="p">,</span> <span class="n">out_proj_weight</span><span class="p">,</span> <span class="n">out_proj_bias</span><span class="p">)</span>
        <span class="n">attn_output</span> <span class="o">=</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">tgt_len</span><span class="p">,</span> <span class="n">bsz</span><span class="p">,</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">size</span><span class="p">(</span><span class="mi">1</span><span class="p">))</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">is_batched</span><span class="p">:</span>
            <span class="c1"># squeeze the output if input was unbatched</span>
            <span class="n">attn_output</span> <span class="o">=</span> <span class="n">attn_output</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">attn_output</span><span class="p">,</span> <span class="kc">None</span>
</pre></div>

             </article>
             
            </div>
            <footer>
  

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