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

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

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

  <!-- Preload the theme fonts -->

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

<!-- Preload the katex fonts -->

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

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

      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

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

          <li>
          <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="resource-option with-down-arrow">
                Edge
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/edge">
                  <span class="dropdown-title">About PyTorch Edge</span>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/executorch">
                  <span class="dropdown-title">ExecuTorch</span>
                </a>
              </div>
            </div>  
          </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.2 &#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">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="../../cpu.html">torch.cpu</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../torch_cuda_memory.html">Understanding CUDA Memory Usage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../torch_cuda_memory.html#generating-a-snapshot">Generating a Snapshot</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../torch_cuda_memory.html#using-the-visualizer">Using the visualizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../torch_cuda_memory.html#snapshot-api-reference">Snapshot API Reference</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../mps.html">torch.mps</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../backends.html">torch.backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../export.html">torch.export</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../distributed.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="../../torch.compiler.html">torch.compiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../fft.html">torch.fft</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../func.html">torch.func</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../futures.html">torch.futures</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../fx.html">torch.fx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../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="../../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="../../utils.html">torch.utils</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../benchmark_utils.html">torch.utils.benchmark</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../deterministic.html">torch.utils.deterministic</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../jit_utils.html">torch.utils.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../mobile_optimizer.html">torch.utils.mobile_optimizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../config_mod.html">torch.__config__</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../logging.html">torch._logging</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.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">

        
          <!-- Google Tag Manager (noscript) -->
          <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-T8XT4PS"
          height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>
          <!-- End Google Tag Manager (noscript) -->
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <h1>Source code for torch.functional</h1><div class="highlight"><pre>
<span></span><span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">List</span><span class="p">,</span> <span class="n">Tuple</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">,</span> <span class="n">TYPE_CHECKING</span>
<span class="p">)</span>
<span class="kn">import</span> <span class="nn">operator</span>
<span class="kn">import</span> <span class="nn">itertools</span>

<span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">from</span> <span class="nn">torch._C</span> <span class="kn">import</span> <span class="n">_add_docstr</span>
<span class="kn">import</span> <span class="nn">torch.nn.functional</span> <span class="k">as</span> <span class="nn">F</span>
<span class="kn">from</span> <span class="nn">._lowrank</span> <span class="kn">import</span> <span class="n">svd_lowrank</span><span class="p">,</span> <span class="n">pca_lowrank</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">._jit_internal</span> <span class="kn">import</span> <span class="n">boolean_dispatch</span>
<span class="kn">from</span> <span class="nn">._jit_internal</span> <span class="kn">import</span> <span class="n">_overload</span> <span class="k">as</span> <span class="n">overload</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="kn">from</span> <span class="nn">torch</span> <span class="kn">import</span> <span class="n">_VF</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span>
    <span class="s1">&#39;atleast_1d&#39;</span><span class="p">,</span>
    <span class="s1">&#39;atleast_2d&#39;</span><span class="p">,</span>
    <span class="s1">&#39;atleast_3d&#39;</span><span class="p">,</span>
    <span class="s1">&#39;align_tensors&#39;</span><span class="p">,</span>
    <span class="s1">&#39;broadcast_shapes&#39;</span><span class="p">,</span>
    <span class="s1">&#39;broadcast_tensors&#39;</span><span class="p">,</span>
    <span class="s1">&#39;cartesian_prod&#39;</span><span class="p">,</span>
    <span class="s1">&#39;block_diag&#39;</span><span class="p">,</span>
    <span class="s1">&#39;cdist&#39;</span><span class="p">,</span>
    <span class="s1">&#39;chain_matmul&#39;</span><span class="p">,</span>
    <span class="s1">&#39;einsum&#39;</span><span class="p">,</span>
    <span class="s1">&#39;istft&#39;</span><span class="p">,</span>
    <span class="s1">&#39;lu&#39;</span><span class="p">,</span>
    <span class="s1">&#39;norm&#39;</span><span class="p">,</span>
    <span class="s1">&#39;meshgrid&#39;</span><span class="p">,</span>
    <span class="s1">&#39;pca_lowrank&#39;</span><span class="p">,</span>
    <span class="s1">&#39;split&#39;</span><span class="p">,</span>
    <span class="s1">&#39;stft&#39;</span><span class="p">,</span>
    <span class="s1">&#39;svd_lowrank&#39;</span><span class="p">,</span>
    <span class="s1">&#39;tensordot&#39;</span><span class="p">,</span>
    <span class="s1">&#39;unique&#39;</span><span class="p">,</span>
    <span class="s1">&#39;unique_consecutive&#39;</span><span class="p">,</span>
    <span class="s1">&#39;unravel_index&#39;</span><span class="p">,</span>
<span class="p">]</span>


<div class="viewcode-block" id="broadcast_tensors"><a class="viewcode-back" href="../../generated/torch.broadcast_tensors.html#torch.broadcast_tensors">[docs]</a><span class="k">def</span> <span class="nf">broadcast_tensors</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;broadcast_tensors(*tensors) -&gt; List of Tensors</span>

<span class="sd">    Broadcasts the given tensors according to :ref:`broadcasting-semantics`.</span>

<span class="sd">    Args:</span>
<span class="sd">        *tensors: any number of tensors of the same type</span>

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

<span class="sd">        More than one element of a broadcasted 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 tensors, please clone them first.</span>

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

<span class="sd">        &gt;&gt;&gt; x = torch.arange(3).view(1, 3)</span>
<span class="sd">        &gt;&gt;&gt; y = torch.arange(2).view(2, 1)</span>
<span class="sd">        &gt;&gt;&gt; a, b = torch.broadcast_tensors(x, y)</span>
<span class="sd">        &gt;&gt;&gt; a.size()</span>
<span class="sd">        torch.Size([2, 3])</span>
<span class="sd">        &gt;&gt;&gt; a</span>
<span class="sd">        tensor([[0, 1, 2],</span>
<span class="sd">                [0, 1, 2]])</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">broadcast_tensors</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">broadcast_tensors</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<div class="viewcode-block" id="broadcast_shapes"><a class="viewcode-back" href="../../generated/torch.broadcast_shapes.html#torch.broadcast_shapes">[docs]</a><span class="k">def</span> <span class="nf">broadcast_shapes</span><span class="p">(</span><span class="o">*</span><span class="n">shapes</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;broadcast_shapes(*shapes) -&gt; Size</span>

<span class="sd">    Similar to :func:`broadcast_tensors` but for shapes.</span>

<span class="sd">    This is equivalent to</span>
<span class="sd">    ``torch.broadcast_tensors(*map(torch.empty, shapes))[0].shape``</span>
<span class="sd">    but avoids the need create to intermediate tensors. This is useful for</span>
<span class="sd">    broadcasting tensors of common batch shape but different rightmost shape,</span>
<span class="sd">    e.g. to broadcast mean vectors with covariance matrices.</span>

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

<span class="sd">        &gt;&gt;&gt; torch.broadcast_shapes((2,), (3, 1), (1, 1, 1))</span>
<span class="sd">        torch.Size([1, 3, 2])</span>

<span class="sd">    Args:</span>
<span class="sd">        \*shapes (torch.Size): Shapes of tensors.</span>

<span class="sd">    Returns:</span>
<span class="sd">        shape (torch.Size): A shape compatible with all input shapes.</span>

<span class="sd">    Raises:</span>
<span class="sd">        RuntimeError: If shapes are incompatible.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="c1"># TODO Move this to C++ once the jit has better support for torch.Size.</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_tracing</span><span class="p">():</span>
        <span class="n">max_len</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">shape</span> <span class="ow">in</span> <span class="n">shapes</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
                <span class="k">if</span> <span class="n">max_len</span> <span class="o">&lt;</span> <span class="mi">1</span><span class="p">:</span>
                    <span class="n">max_len</span> <span class="o">=</span> <span class="mi">1</span>
            <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
                <span class="n">s</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">max_len</span> <span class="o">&lt;</span> <span class="n">s</span><span class="p">:</span>
                    <span class="n">max_len</span> <span class="o">=</span> <span class="n">s</span>
        <span class="n">result</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_len</span>
        <span class="k">for</span> <span class="n">shape</span> <span class="ow">in</span> <span class="n">shapes</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
                <span class="n">shape</span> <span class="o">=</span> <span class="p">(</span><span class="n">shape</span><span class="p">,)</span>
            <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</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="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span> <span class="o">-</span> <span class="nb">len</span><span class="p">(</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="k">if</span> <span class="n">shape</span><span class="p">[</span><span class="n">i</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">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Trying to create tensor with negative dimension (</span><span class="si">{</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="si">}</span><span class="s2">): (</span><span class="si">{</span><span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="si">}</span><span class="s2">)&quot;</span><span class="p">)</span>
                    <span class="k">if</span> <span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">==</span> <span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">]:</span>
                        <span class="k">continue</span>
                    <span class="k">if</span> <span class="n">result</span><span class="p">[</span><span class="n">i</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">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Shape mismatch: objects cannot be broadcast to a single shape&quot;</span><span class="p">)</span>
                    <span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Input shapes should be of type ints, a tuple of ints, or a list of ints, got &quot;</span><span class="p">,</span> <span class="n">shape</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">Size</span><span class="p">(</span><span class="n">result</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># with implementation above, torch.jit.trace hardcodes the sizes which makes subsequent replays fail</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">scalar</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">zeros</span><span class="p">((),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span><span class="p">)</span>
            <span class="n">tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">scalar</span><span class="o">.</span><span class="n">expand</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span> <span class="k">for</span> <span class="n">shape</span> <span class="ow">in</span> <span class="n">shapes</span><span class="p">]</span>
            <span class="n">tensors</span> <span class="o">=</span> <span class="n">broadcast_tensors</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span></div>


<div class="viewcode-block" id="split"><a class="viewcode-back" href="../../generated/torch.split.html#torch.split">[docs]</a><span class="k">def</span> <span class="nf">split</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">split_size_or_sections</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="n">List</span><span class="p">[</span><span class="nb">int</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">0</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="o">...</span><span class="p">]:</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Splits the tensor into chunks. Each chunk is a view of the original tensor.</span>

<span class="sd">    If :attr:`split_size_or_sections` is an integer type, then :attr:`tensor` will</span>
<span class="sd">    be split into equally sized chunks (if possible). Last chunk will be smaller if</span>
<span class="sd">    the tensor size along the given dimension :attr:`dim` is not divisible by</span>
<span class="sd">    :attr:`split_size`.</span>

<span class="sd">    If :attr:`split_size_or_sections` is a list, then :attr:`tensor` will be split</span>
<span class="sd">    into ``len(split_size_or_sections)`` chunks with sizes in :attr:`dim` according</span>
<span class="sd">    to :attr:`split_size_or_sections`.</span>

<span class="sd">    Args:</span>
<span class="sd">        tensor (Tensor): tensor to split.</span>
<span class="sd">        split_size_or_sections (int) or (list(int)): size of a single chunk or</span>
<span class="sd">            list of sizes for each chunk</span>
<span class="sd">        dim (int): dimension along which to split the tensor.</span>

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

<span class="sd">        &gt;&gt;&gt; a = torch.arange(10).reshape(5, 2)</span>
<span class="sd">        &gt;&gt;&gt; a</span>
<span class="sd">        tensor([[0, 1],</span>
<span class="sd">                [2, 3],</span>
<span class="sd">                [4, 5],</span>
<span class="sd">                [6, 7],</span>
<span class="sd">                [8, 9]])</span>
<span class="sd">        &gt;&gt;&gt; torch.split(a, 2)</span>
<span class="sd">        (tensor([[0, 1],</span>
<span class="sd">                 [2, 3]]),</span>
<span class="sd">         tensor([[4, 5],</span>
<span class="sd">                 [6, 7]]),</span>
<span class="sd">         tensor([[8, 9]]))</span>
<span class="sd">        &gt;&gt;&gt; torch.split(a, [1, 4])</span>
<span class="sd">        (tensor([[0, 1]]),</span>
<span class="sd">         tensor([[2, 3],</span>
<span class="sd">                 [4, 5],</span>
<span class="sd">                 [6, 7],</span>
<span class="sd">                 [8, 9]]))</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">tensor</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">split</span><span class="p">,</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">split_size_or_sections</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="c1"># Overwriting reason:</span>
    <span class="c1"># This dispatches to two ATen functions depending on the type of</span>
    <span class="c1"># split_size_or_sections. The branching code is in _tensor.py, which we</span>
    <span class="c1"># call here.</span>
    <span class="k">return</span> <span class="n">tensor</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">split_size_or_sections</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span></div>


<div class="viewcode-block" id="einsum"><a class="viewcode-back" href="../../generated/torch.einsum.html#torch.einsum">[docs]</a><span class="k">def</span> <span class="nf">einsum</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;einsum(equation, *operands) -&gt; Tensor</span>

<span class="sd">    Sums the product of the elements of the input :attr:`operands` along dimensions specified using a notation</span>
<span class="sd">    based on the Einstein summation convention.</span>

<span class="sd">    Einsum allows computing many common multi-dimensional linear algebraic array operations by representing them</span>
<span class="sd">    in a short-hand format based on the Einstein summation convention, given by :attr:`equation`. The details of</span>
<span class="sd">    this format are described below, but the general idea is to label every dimension of the input :attr:`operands`</span>
<span class="sd">    with some subscript and define which subscripts are part of the output. The output is then computed by summing</span>
<span class="sd">    the product of the elements of the :attr:`operands` along the dimensions whose subscripts are not part of the</span>
<span class="sd">    output. For example, matrix multiplication can be computed using einsum as `torch.einsum(&quot;ij,jk-&gt;ik&quot;, A, B)`.</span>
<span class="sd">    Here, j is the summation subscript and i and k the output subscripts (see section below for more details on why).</span>

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

<span class="sd">        The :attr:`equation` string specifies the subscripts (letters in `[a-zA-Z]`) for each dimension of</span>
<span class="sd">        the input :attr:`operands` in the same order as the dimensions, separating subscripts for each operand by a</span>
<span class="sd">        comma (&#39;,&#39;), e.g. `&#39;ij,jk&#39;` specify subscripts for two 2D operands. The dimensions labeled with the same subscript</span>
<span class="sd">        must be broadcastable, that is, their size must either match or be `1`. The exception is if a subscript is</span>
<span class="sd">        repeated for the same input operand, in which case the dimensions labeled with this subscript for this operand</span>
<span class="sd">        must match in size and the operand will be replaced by its diagonal along these dimensions. The subscripts that</span>
<span class="sd">        appear exactly once in the :attr:`equation` will be part of the output, sorted in increasing alphabetical order.</span>
<span class="sd">        The output is computed by multiplying the input :attr:`operands` element-wise, with their dimensions aligned based</span>
<span class="sd">        on the subscripts, and then summing out the dimensions whose subscripts are not part of the output.</span>

<span class="sd">        Optionally, the output subscripts can be explicitly defined by adding an arrow (&#39;-&gt;&#39;) at the end of the equation</span>
<span class="sd">        followed by the subscripts for the output. For instance, the following equation computes the transpose of a</span>
<span class="sd">        matrix multiplication: &#39;ij,jk-&gt;ki&#39;. The output subscripts must appear at least once for some input operand and</span>
<span class="sd">        at most once for the output.</span>

<span class="sd">        Ellipsis (&#39;...&#39;) can be used in place of subscripts to broadcast the dimensions covered by the ellipsis.</span>
<span class="sd">        Each input operand may contain at most one ellipsis which will cover the dimensions not covered by subscripts,</span>
<span class="sd">        e.g. for an input operand with 5 dimensions, the ellipsis in the equation `&#39;ab...c&#39;` cover the third and fourth</span>
<span class="sd">        dimensions. The ellipsis does not need to cover the same number of dimensions across the :attr:`operands` but the</span>
<span class="sd">        &#39;shape&#39; of the ellipsis (the size of the dimensions covered by them) must broadcast together. If the output is not</span>
<span class="sd">        explicitly defined with the arrow (&#39;-&gt;&#39;) notation, the ellipsis will come first in the output (left-most dimensions),</span>
<span class="sd">        before the subscript labels that appear exactly once for the input operands. e.g. the following equation implements</span>
<span class="sd">        batch matrix multiplication `&#39;...ij,...jk&#39;`.</span>

<span class="sd">        A few final notes: the equation may contain whitespaces between the different elements (subscripts, ellipsis,</span>
<span class="sd">        arrow and comma) but something like `&#39;. . .&#39;` is not valid. An empty string `&#39;&#39;` is valid for scalar operands.</span>

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

<span class="sd">        ``torch.einsum`` handles ellipsis (&#39;...&#39;) differently from NumPy in that it allows dimensions</span>
<span class="sd">        covered by the ellipsis to be summed over, that is, ellipsis are not required to be part of the output.</span>

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

<span class="sd">        This function uses opt_einsum (https://optimized-einsum.readthedocs.io/en/stable/) to speed up computation or to</span>
<span class="sd">        consume less memory by optimizing contraction order. This optimization occurs when there are at least three</span>
<span class="sd">        inputs, since the order does not matter otherwise. Note that finding _the_ optimal path is an NP-hard problem,</span>
<span class="sd">        thus, opt_einsum relies on different heuristics to achieve near-optimal results. If opt_einsum is not available,</span>
<span class="sd">        the default order is to contract from left to right.</span>

<span class="sd">        To bypass this default behavior, add the following line to disable the usage of opt_einsum and skip path</span>
<span class="sd">        calculation: `torch.backends.opt_einsum.enabled = False`</span>

<span class="sd">        To specify which strategy you&#39;d like for opt_einsum to compute the contraction path, add the following line:</span>
<span class="sd">        `torch.backends.opt_einsum.strategy = &#39;auto&#39;`. The default strategy is &#39;auto&#39;, and we also support &#39;greedy&#39; and</span>
<span class="sd">        &#39;optimal&#39;. Disclaimer that the runtime of &#39;optimal&#39; is factorial in the number of inputs! See more details in</span>
<span class="sd">        the opt_einsum documentation (https://optimized-einsum.readthedocs.io/en/stable/path_finding.html).</span>

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

<span class="sd">        As of PyTorch 1.10 :func:`torch.einsum` also supports the sublist format (see examples below). In this format,</span>
<span class="sd">        subscripts for each operand are specified by sublists, list of integers in the range [0, 52). These sublists</span>
<span class="sd">        follow their operands, and an extra sublist can appear at the end of the input to specify the output&#39;s</span>
<span class="sd">        subscripts., e.g. `torch.einsum(op1, sublist1, op2, sublist2, ..., [subslist_out])`. Python&#39;s `Ellipsis` object</span>
<span class="sd">        may be provided in a sublist to enable broadcasting as described in the Equation section above.</span>

<span class="sd">    Args:</span>
<span class="sd">        equation (str): The subscripts for the Einstein summation.</span>
<span class="sd">        operands (List[Tensor]): The tensors to compute the Einstein summation of.</span>

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

<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; # trace</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(&#39;ii&#39;, torch.randn(4, 4))</span>
<span class="sd">        tensor(-1.2104)</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; # diagonal</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(&#39;ii-&gt;i&#39;, torch.randn(4, 4))</span>
<span class="sd">        tensor([-0.1034,  0.7952, -0.2433,  0.4545])</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; # outer product</span>
<span class="sd">        &gt;&gt;&gt; x = torch.randn(5)</span>
<span class="sd">        &gt;&gt;&gt; y = torch.randn(4)</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(&#39;i,j-&gt;ij&#39;, x, y)</span>
<span class="sd">        tensor([[ 0.1156, -0.2897, -0.3918,  0.4963],</span>
<span class="sd">                [-0.3744,  0.9381,  1.2685, -1.6070],</span>
<span class="sd">                [ 0.7208, -1.8058, -2.4419,  3.0936],</span>
<span class="sd">                [ 0.1713, -0.4291, -0.5802,  0.7350],</span>
<span class="sd">                [ 0.5704, -1.4290, -1.9323,  2.4480]])</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; # batch matrix multiplication</span>
<span class="sd">        &gt;&gt;&gt; As = torch.randn(3, 2, 5)</span>
<span class="sd">        &gt;&gt;&gt; Bs = torch.randn(3, 5, 4)</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(&#39;bij,bjk-&gt;bik&#39;, As, Bs)</span>
<span class="sd">        tensor([[[-1.0564, -1.5904,  3.2023,  3.1271],</span>
<span class="sd">                [-1.6706, -0.8097, -0.8025, -2.1183]],</span>

<span class="sd">                [[ 4.2239,  0.3107, -0.5756, -0.2354],</span>
<span class="sd">                [-1.4558, -0.3460,  1.5087, -0.8530]],</span>

<span class="sd">                [[ 2.8153,  1.8787, -4.3839, -1.2112],</span>
<span class="sd">                [ 0.3728, -2.1131,  0.0921,  0.8305]]])</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; # with sublist format and ellipsis</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(As, [..., 0, 1], Bs, [..., 1, 2], [..., 0, 2])</span>
<span class="sd">        tensor([[[-1.0564, -1.5904,  3.2023,  3.1271],</span>
<span class="sd">                [-1.6706, -0.8097, -0.8025, -2.1183]],</span>

<span class="sd">                [[ 4.2239,  0.3107, -0.5756, -0.2354],</span>
<span class="sd">                [-1.4558, -0.3460,  1.5087, -0.8530]],</span>

<span class="sd">                [[ 2.8153,  1.8787, -4.3839, -1.2112],</span>
<span class="sd">                [ 0.3728, -2.1131,  0.0921,  0.8305]]])</span>

<span class="sd">        &gt;&gt;&gt; # batch permute</span>
<span class="sd">        &gt;&gt;&gt; A = torch.randn(2, 3, 4, 5)</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(&#39;...ij-&gt;...ji&#39;, A).shape</span>
<span class="sd">        torch.Size([2, 3, 5, 4])</span>

<span class="sd">        &gt;&gt;&gt; # equivalent to torch.nn.functional.bilinear</span>
<span class="sd">        &gt;&gt;&gt; A = torch.randn(3, 5, 4)</span>
<span class="sd">        &gt;&gt;&gt; l = torch.randn(2, 5)</span>
<span class="sd">        &gt;&gt;&gt; r = torch.randn(2, 4)</span>
<span class="sd">        &gt;&gt;&gt; torch.einsum(&#39;bn,anm,bm-&gt;ba&#39;, l, A, r)</span>
<span class="sd">        tensor([[-0.3430, -5.2405,  0.4494],</span>
<span class="sd">                [ 0.3311,  5.5201, -3.0356]])</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="kn">import</span> <span class="nn">torch.backends.opt_einsum</span> <span class="k">as</span> <span class="nn">opt_einsum</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">args</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">ValueError</span><span class="p">(</span><span class="s1">&#39;einsum(): must specify the equation string and at least one operand, &#39;</span>
                         <span class="s1">&#39;or at least one operand and its subscripts list&#39;</span><span class="p">)</span>

    <span class="n">equation</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">operands</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">args</span><span class="p">[</span><span class="mi">0</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"># Convert the subscript list format which is an interleaving of operand and its subscripts</span>
        <span class="c1"># list with an optional output subscripts list at the end (see documentation for more details on this)</span>
        <span class="c1"># to the equation string format by creating the equation string from the subscripts list and grouping the</span>
        <span class="c1"># input operands into a tensorlist (List[Tensor]).</span>
        <span class="k">def</span> <span class="nf">parse_subscript</span><span class="p">(</span><span class="n">n</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">n</span> <span class="o">==</span> <span class="bp">Ellipsis</span><span class="p">:</span>
                <span class="k">return</span> <span class="s1">&#39;...&#39;</span>
            <span class="k">if</span> <span class="n">n</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">n</span> <span class="o">&lt;</span> <span class="mi">26</span><span class="p">:</span>
                <span class="k">return</span> <span class="nb">chr</span><span class="p">(</span><span class="nb">ord</span><span class="p">(</span><span class="s1">&#39;A&#39;</span><span class="p">)</span> <span class="o">+</span> <span class="n">n</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">n</span> <span class="o">&gt;=</span> <span class="mi">26</span> <span class="ow">and</span> <span class="n">n</span> <span class="o">&lt;</span> <span class="mi">52</span><span class="p">:</span>
                <span class="k">return</span> <span class="nb">chr</span><span class="p">(</span><span class="nb">ord</span><span class="p">(</span><span class="s1">&#39;a&#39;</span><span class="p">)</span> <span class="o">+</span> <span class="n">n</span> <span class="o">-</span> <span class="mi">26</span><span class="p">)</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;einsum(): subscript in subscript list is not within the valid range [0, 52)&#39;</span><span class="p">)</span>

        <span class="c1"># Parse subscripts for input operands</span>
        <span class="n">equation</span> <span class="o">=</span> <span class="s1">&#39;,&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="s1">&#39;&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">parse_subscript</span><span class="p">(</span><span class="n">s</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">l</span><span class="p">)</span> <span class="k">for</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">::</span><span class="mi">2</span><span class="p">])</span>

        <span class="c1"># Parse optional output subscripts (provided when the number of arguments is odd)</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">args</span><span class="p">)</span> <span class="o">%</span> <span class="mi">2</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">equation</span> <span class="o">+=</span> <span class="s1">&#39;-&gt;&#39;</span> <span class="o">+</span> <span class="s1">&#39;&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">parse_subscript</span><span class="p">(</span><span class="n">s</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">args</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
            <span class="n">operands</span> <span class="o">=</span> <span class="n">args</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">:</span><span class="mi">2</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">operands</span> <span class="o">=</span> <span class="n">args</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="n">equation</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">operands</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>

    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">operands</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">einsum</span><span class="p">,</span> <span class="n">operands</span><span class="p">,</span> <span class="n">equation</span><span class="p">,</span> <span class="o">*</span><span class="n">operands</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">operands</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">operands</span><span class="p">[</span><span class="mi">0</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="c1"># the old interface of passing the operands as one list argument</span>
        <span class="n">_operands</span> <span class="o">=</span> <span class="n">operands</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="c1"># recurse incase operands contains value that has torch function</span>
        <span class="c1"># in the original implementation this line is omitted</span>
        <span class="k">return</span> <span class="n">einsum</span><span class="p">(</span><span class="n">equation</span><span class="p">,</span> <span class="o">*</span><span class="n">_operands</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">operands</span><span class="p">)</span> <span class="o">&lt;=</span> <span class="mi">2</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">opt_einsum</span><span class="o">.</span><span class="n">enabled</span><span class="p">:</span>
        <span class="c1"># the path for contracting 0 or 1 time(s) is already optimized</span>
        <span class="c1"># or the user has disabled using opt_einsum</span>
        <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">einsum</span><span class="p">(</span><span class="n">equation</span><span class="p">,</span> <span class="n">operands</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>

    <span class="n">path</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">if</span> <span class="n">opt_einsum</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
        <span class="n">_opt_einsum</span> <span class="o">=</span> <span class="n">opt_einsum</span><span class="o">.</span><span class="n">get_opt_einsum</span><span class="p">()</span>
        <span class="n">tupled_path</span> <span class="o">=</span> <span class="n">_opt_einsum</span><span class="o">.</span><span class="n">contract_path</span><span class="p">(</span><span class="n">equation</span><span class="p">,</span> <span class="o">*</span><span class="n">operands</span><span class="p">,</span> <span class="n">optimize</span><span class="o">=</span><span class="n">opt_einsum</span><span class="o">.</span><span class="n">strategy</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
        <span class="c1"># flatten path for dispatching to C++</span>
        <span class="n">path</span> <span class="o">=</span> <span class="p">[</span><span class="n">item</span> <span class="k">for</span> <span class="n">pair</span> <span class="ow">in</span> <span class="n">tupled_path</span> <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">pair</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">einsum</span><span class="p">(</span><span class="n">equation</span><span class="p">,</span> <span class="n">operands</span><span class="p">,</span> <span class="n">path</span><span class="o">=</span><span class="n">path</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<span class="c1"># This wrapper exists to support variadic args.</span>
<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="c1"># The JIT doesn&#39;t understand Union, so only add type annotation for mypy</span>
    <span class="k">def</span> <span class="nf">meshgrid</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]],</span>
                 <span class="n">indexing</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="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="o">...</span><span class="p">]:</span>
        <span class="k">return</span> <span class="n">_meshgrid</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">,</span> <span class="n">indexing</span><span class="o">=</span><span class="n">indexing</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<div class="viewcode-block" id="meshgrid"><a class="viewcode-back" href="../../generated/torch.meshgrid.html#torch.meshgrid">[docs]</a>    <span class="k">def</span> <span class="nf">meshgrid</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">,</span> <span class="n">indexing</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="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="o">...</span><span class="p">]:</span>
<span class="w">        </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Creates grids of coordinates specified by the 1D inputs in `attr`:tensors.</span>

<span class="sd">        This is helpful when you want to visualize data over some</span>
<span class="sd">        range of inputs. See below for a plotting example.</span>

<span class="sd">        Given :math:`N` 1D tensors :math:`T_0 \ldots T_{N-1}` as</span>
<span class="sd">        inputs with corresponding sizes :math:`S_0 \ldots S_{N-1}`,</span>
<span class="sd">        this creates :math:`N` N-dimensional tensors :math:`G_0 \ldots</span>
<span class="sd">        G_{N-1}`, each with shape :math:`(S_0, ..., S_{N-1})` where</span>
<span class="sd">        the output :math:`G_i` is constructed by expanding :math:`T_i`</span>
<span class="sd">        to the result shape.</span>

<span class="sd">        .. note::</span>
<span class="sd">            0D inputs are treated equivalently to 1D inputs of a</span>
<span class="sd">            single element.</span>

<span class="sd">        .. warning::</span>
<span class="sd">            `torch.meshgrid(*tensors)` currently has the same behavior</span>
<span class="sd">            as calling `numpy.meshgrid(*arrays, indexing=&#39;ij&#39;)`.</span>

<span class="sd">            In the future `torch.meshgrid` will transition to</span>
<span class="sd">            `indexing=&#39;xy&#39;` as the default.</span>

<span class="sd">            https://github.com/pytorch/pytorch/issues/50276 tracks</span>
<span class="sd">            this issue with the goal of migrating to NumPy&#39;s behavior.</span>

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

<span class="sd">            :func:`torch.cartesian_prod` has the same effect but it</span>
<span class="sd">            collects the data in a tensor of vectors.</span>

<span class="sd">        Args:</span>
<span class="sd">            tensors (list of Tensor): list of scalars or 1 dimensional tensors. Scalars will be</span>
<span class="sd">                treated as tensors of size :math:`(1,)` automatically</span>

<span class="sd">            indexing: (str, optional): the indexing mode, either &quot;xy&quot;</span>
<span class="sd">                or &quot;ij&quot;, defaults to &quot;ij&quot;. See warning for future changes.</span>

<span class="sd">                If &quot;xy&quot; is selected, the first dimension corresponds</span>
<span class="sd">                to the cardinality of the second input and the second</span>
<span class="sd">                dimension corresponds to the cardinality of the first</span>
<span class="sd">                input.</span>

<span class="sd">                If &quot;ij&quot; is selected, the dimensions are in the same</span>
<span class="sd">                order as the cardinality of the inputs.</span>

<span class="sd">        Returns:</span>
<span class="sd">            seq (sequence of Tensors): If the input has :math:`N`</span>
<span class="sd">            tensors of size :math:`S_0 \ldots S_{N-1}``, then the</span>
<span class="sd">            output will also have :math:`N` tensors, where each tensor</span>
<span class="sd">            is of shape :math:`(S_0, ..., S_{N-1})`.</span>

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

<span class="sd">            &gt;&gt;&gt; x = torch.tensor([1, 2, 3])</span>
<span class="sd">            &gt;&gt;&gt; y = torch.tensor([4, 5, 6])</span>

<span class="sd">            Observe the element-wise pairings across the grid, (1, 4),</span>
<span class="sd">            (1, 5), ..., (3, 6). This is the same thing as the</span>
<span class="sd">            cartesian product.</span>
<span class="sd">            &gt;&gt;&gt; grid_x, grid_y = torch.meshgrid(x, y, indexing=&#39;ij&#39;)</span>
<span class="sd">            &gt;&gt;&gt; grid_x</span>
<span class="sd">            tensor([[1, 1, 1],</span>
<span class="sd">                    [2, 2, 2],</span>
<span class="sd">                    [3, 3, 3]])</span>
<span class="sd">            &gt;&gt;&gt; grid_y</span>
<span class="sd">            tensor([[4, 5, 6],</span>
<span class="sd">                    [4, 5, 6],</span>
<span class="sd">                    [4, 5, 6]])</span>

<span class="sd">            This correspondence can be seen when these grids are</span>
<span class="sd">            stacked properly.</span>
<span class="sd">            &gt;&gt;&gt; torch.equal(torch.cat(tuple(torch.dstack([grid_x, grid_y]))),</span>
<span class="sd">            ...             torch.cartesian_prod(x, y))</span>
<span class="sd">            True</span>

<span class="sd">            `torch.meshgrid` is commonly used to produce a grid for</span>
<span class="sd">            plotting.</span>
<span class="sd">            &gt;&gt;&gt; # xdoctest: +REQUIRES(module:matplotlib)</span>
<span class="sd">            &gt;&gt;&gt; # xdoctest: +REQUIRES(env:DOCTEST_SHOW)</span>
<span class="sd">            &gt;&gt;&gt; import matplotlib.pyplot as plt</span>
<span class="sd">            &gt;&gt;&gt; xs = torch.linspace(-5, 5, steps=100)</span>
<span class="sd">            &gt;&gt;&gt; ys = torch.linspace(-5, 5, steps=100)</span>
<span class="sd">            &gt;&gt;&gt; x, y = torch.meshgrid(xs, ys, indexing=&#39;xy&#39;)</span>
<span class="sd">            &gt;&gt;&gt; z = torch.sin(torch.sqrt(x * x + y * y))</span>
<span class="sd">            &gt;&gt;&gt; ax = plt.axes(projection=&#39;3d&#39;)</span>
<span class="sd">            &gt;&gt;&gt; ax.plot_surface(x.numpy(), y.numpy(), z.numpy())</span>
<span class="sd">            &gt;&gt;&gt; plt.show()</span>

<span class="sd">        .. image:: ../_static/img/meshgrid.png</span>
<span class="sd">            :width: 512</span>

<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">_meshgrid</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">,</span> <span class="n">indexing</span><span class="o">=</span><span class="n">indexing</span><span class="p">)</span></div>


<span class="k">def</span> <span class="nf">_meshgrid</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">,</span> <span class="n">indexing</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]):</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">meshgrid</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</span><span class="p">,</span> <span class="n">indexing</span><span class="o">=</span><span class="n">indexing</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">tensors</span><span class="p">[</span><span class="mi">0</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="c1"># the old interface of passing the operands as one list argument</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>  <span class="c1"># type: ignore[assignment]</span>

    <span class="c1"># Continue allowing call of old method that takes no indexing</span>
    <span class="c1"># kwarg for forward compatibility reasons.</span>
    <span class="c1">#</span>
    <span class="c1"># Remove this two weeks after landing.</span>
    <span class="n">kwargs</span> <span class="o">=</span> <span class="p">{}</span> <span class="k">if</span> <span class="n">indexing</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="p">{</span><span class="s1">&#39;indexing&#39;</span><span class="p">:</span> <span class="n">indexing</span><span class="p">}</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">tensors</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>


<div class="viewcode-block" id="stft"><a class="viewcode-back" href="../../generated/torch.stft.html#torch.stft">[docs]</a><span class="k">def</span> <span class="nf">stft</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">n_fft</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">hop_length</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">win_length</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">window</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">center</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">pad_mode</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s1">&#39;reflect&#39;</span><span class="p">,</span> <span class="n">normalized</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">onesided</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">return_complex</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="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Short-time Fourier transform (STFT).</span>

<span class="sd">    .. warning::</span>
<span class="sd">        From version 1.8.0, :attr:`return_complex` must always be given</span>
<span class="sd">        explicitly for real inputs and `return_complex=False` has been</span>
<span class="sd">        deprecated. Strongly prefer `return_complex=True` as in a future</span>
<span class="sd">        pytorch release, this function will only return complex tensors.</span>

<span class="sd">        Note that :func:`torch.view_as_real` can be used to recover a real</span>
<span class="sd">        tensor with an extra last dimension for real and imaginary components.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        From version 2.1, a warning will be provided if a :attr:`window` is</span>
<span class="sd">        not specified. In a future release, this attribute will be required.</span>
<span class="sd">        Not providing a window currently defaults to using a rectangular window,</span>
<span class="sd">        which may result in undesirable artifacts. Consider using tapered windows,</span>
<span class="sd">        such as :func:`torch.hann_window`.</span>

<span class="sd">    The STFT computes the Fourier transform of short overlapping windows of the</span>
<span class="sd">    input. This giving frequency components of the signal as they change over</span>
<span class="sd">    time. The interface of this function is modeled after (but *not* a drop-in</span>
<span class="sd">    replacement for) librosa_ stft function.</span>

<span class="sd">    .. _librosa: https://librosa.org/doc/latest/generated/librosa.stft.html</span>

<span class="sd">    Ignoring the optional batch dimension, this method computes the following</span>
<span class="sd">    expression:</span>

<span class="sd">    .. math::</span>
<span class="sd">        X[\omega, m] = \sum_{k = 0}^{\text{win\_length-1}}%</span>
<span class="sd">                            \text{window}[k]\ \text{input}[m \times \text{hop\_length} + k]\ %</span>
<span class="sd">                            \exp\left(- j \frac{2 \pi \cdot \omega k}{\text{n\_fft}}\right),</span>

<span class="sd">    where :math:`m` is the index of the sliding window, and :math:`\omega` is</span>
<span class="sd">    the frequency :math:`0 \leq \omega &lt; \text{n\_fft}` for ``onesided=False``,</span>
<span class="sd">    or :math:`0 \leq \omega &lt; \lfloor \text{n\_fft} / 2 \rfloor + 1` for ``onesided=True``.</span>

<span class="sd">    * :attr:`input` must be either a 1-D time sequence or a 2-D batch of time</span>
<span class="sd">      sequences.</span>

<span class="sd">    * If :attr:`hop_length` is ``None`` (default), it is treated as equal to</span>
<span class="sd">      ``floor(n_fft / 4)``.</span>

<span class="sd">    * If :attr:`win_length` is ``None`` (default), it is treated as equal to</span>
<span class="sd">      :attr:`n_fft`.</span>

<span class="sd">    * :attr:`window` can be a 1-D tensor of size :attr:`win_length`, e.g., from</span>
<span class="sd">      :meth:`torch.hann_window`. If :attr:`window` is ``None`` (default), it is</span>
<span class="sd">      treated as if having :math:`1` everywhere in the window. If</span>
<span class="sd">      :math:`\text{win\_length} &lt; \text{n\_fft}`, :attr:`window` will be padded on</span>
<span class="sd">      both sides to length :attr:`n_fft` before being applied.</span>

<span class="sd">    * If :attr:`center` is ``True`` (default), :attr:`input` will be padded on</span>
<span class="sd">      both sides so that the :math:`t`-th frame is centered at time</span>
<span class="sd">      :math:`t \times \text{hop\_length}`. Otherwise, the :math:`t`-th frame</span>
<span class="sd">      begins at time  :math:`t \times \text{hop\_length}`.</span>

<span class="sd">    * :attr:`pad_mode` determines the padding method used on :attr:`input` when</span>
<span class="sd">      :attr:`center` is ``True``. See :meth:`torch.nn.functional.pad` for</span>
<span class="sd">      all available options. Default is ``&quot;reflect&quot;``.</span>

<span class="sd">    * If :attr:`onesided` is ``True`` (default for real input), only values for</span>
<span class="sd">      :math:`\omega` in :math:`\left[0, 1, 2, \dots, \left\lfloor</span>
<span class="sd">      \frac{\text{n\_fft}}{2} \right\rfloor + 1\right]` are returned because</span>
<span class="sd">      the real-to-complex Fourier transform satisfies the conjugate symmetry,</span>
<span class="sd">      i.e., :math:`X[m, \omega] = X[m, \text{n\_fft} - \omega]^*`.</span>
<span class="sd">      Note if the input or window tensors are complex, then :attr:`onesided`</span>
<span class="sd">      output is not possible.</span>

<span class="sd">    * If :attr:`normalized` is ``True`` (default is ``False``), the function</span>
<span class="sd">      returns the normalized STFT results, i.e., multiplied by :math:`(\text{frame\_length})^{-0.5}`.</span>

<span class="sd">    * If :attr:`return_complex` is ``True`` (default if input is complex), the</span>
<span class="sd">      return is a ``input.dim() + 1`` dimensional complex tensor. If ``False``,</span>
<span class="sd">      the output is a ``input.dim() + 2`` dimensional real tensor where the last</span>
<span class="sd">      dimension represents the real and imaginary components.</span>

<span class="sd">    Returns either a complex tensor of size :math:`(* \times N \times T)` if</span>
<span class="sd">    :attr:`return_complex` is true, or a real tensor of size :math:`(* \times N</span>
<span class="sd">    \times T \times 2)`. Where :math:`*` is the optional batch size of</span>
<span class="sd">    :attr:`input`, :math:`N` is the number of frequencies where STFT is applied</span>
<span class="sd">    and :math:`T` is the total number of frames used.</span>

<span class="sd">    .. warning::</span>
<span class="sd">      This function changed signature at version 0.4.1. Calling with the</span>
<span class="sd">      previous signature may cause error or return incorrect result.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): the input tensor of shape `(B?, L)` where `B?` is an optional</span>
<span class="sd">            batch dimension</span>
<span class="sd">        n_fft (int): size of Fourier transform</span>
<span class="sd">        hop_length (int, optional): the distance between neighboring sliding window</span>
<span class="sd">            frames. Default: ``None`` (treated as equal to ``floor(n_fft / 4)``)</span>
<span class="sd">        win_length (int, optional): the size of window frame and STFT filter.</span>
<span class="sd">            Default: ``None``  (treated as equal to :attr:`n_fft`)</span>
<span class="sd">        window (Tensor, optional): the optional window function.</span>
<span class="sd">            Shape must be 1d and `&lt;= n_fft`</span>
<span class="sd">            Default: ``None`` (treated as window of all :math:`1` s)</span>
<span class="sd">        center (bool, optional): whether to pad :attr:`input` on both sides so</span>
<span class="sd">            that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`.</span>
<span class="sd">            Default: ``True``</span>
<span class="sd">        pad_mode (str, optional): controls the padding method used when</span>
<span class="sd">            :attr:`center` is ``True``. Default: ``&quot;reflect&quot;``</span>
<span class="sd">        normalized (bool, optional): controls whether to return the normalized STFT results</span>
<span class="sd">             Default: ``False``</span>
<span class="sd">        onesided (bool, optional): controls whether to return half of results to</span>
<span class="sd">            avoid redundancy for real inputs.</span>
<span class="sd">            Default: ``True`` for real :attr:`input` and :attr:`window`, ``False`` otherwise.</span>
<span class="sd">        return_complex (bool, optional): whether to return a complex tensor, or</span>
<span class="sd">            a real tensor with an extra last dimension for the real and</span>
<span class="sd">            imaginary components.</span>

<span class="sd">            .. versionchanged:: 2.0</span>
<span class="sd">               ``return_complex`` is now a required argument for real inputs,</span>
<span class="sd">               as the default is being transitioned to ``True``.</span>

<span class="sd">            .. deprecated:: 2.0</span>
<span class="sd">               ``return_complex=False`` is deprecated, instead use ``return_complex=True``</span>
<span class="sd">               Note that calling :func:`torch.view_as_real` on the output will</span>
<span class="sd">               recover the deprecated output format.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Tensor: A tensor containing the STFT result with shape `(B?, N, T, C?)` where</span>
<span class="sd">           - `B?` is an optional batch dimnsion from the input</span>
<span class="sd">           - `N` is the number of frequency samples, `(n_fft // 2) + 1` for</span>
<span class="sd">             `onesided=True`, or otherwise `n_fft`.</span>
<span class="sd">           - `T` is the number of frames, `1 + L // hop_length`</span>
<span class="sd">             for `center=True`, or `1 + (L - n_fft) // hop_length` otherwise.</span>
<span class="sd">           - `C?` is an optional length-2 dimension of real and imaginary</span>
<span class="sd">             components, present when `return_complex=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">stft</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">n_fft</span><span class="p">,</span> <span class="n">hop_length</span><span class="o">=</span><span class="n">hop_length</span><span class="p">,</span> <span class="n">win_length</span><span class="o">=</span><span class="n">win_length</span><span class="p">,</span>
            <span class="n">window</span><span class="o">=</span><span class="n">window</span><span class="p">,</span> <span class="n">center</span><span class="o">=</span><span class="n">center</span><span class="p">,</span> <span class="n">pad_mode</span><span class="o">=</span><span class="n">pad_mode</span><span class="p">,</span> <span class="n">normalized</span><span class="o">=</span><span class="n">normalized</span><span class="p">,</span>
            <span class="n">onesided</span><span class="o">=</span><span class="n">onesided</span><span class="p">,</span> <span class="n">return_complex</span><span class="o">=</span><span class="n">return_complex</span><span class="p">)</span>
    <span class="c1"># NOTE: Do not edit. This code will be removed once the forward-compatibility</span>
    <span class="c1">#       period is over for PR #73432</span>
    <span class="k">if</span> <span class="n">center</span><span class="p">:</span>
        <span class="n">signal_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="n">extended_shape</span> <span class="o">=</span> <span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="p">(</span><span class="mi">3</span> <span class="o">-</span> <span class="n">signal_dim</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="n">pad</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">n_fft</span> <span class="o">//</span> <span class="mi">2</span><span class="p">)</span>
        <span class="nb">input</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">pad</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">extended_shape</span><span class="p">),</span> <span class="p">[</span><span class="n">pad</span><span class="p">,</span> <span class="n">pad</span><span class="p">],</span> <span class="n">pad_mode</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">view</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="o">-</span><span class="n">signal_dim</span><span class="p">:])</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">stft</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">n_fft</span><span class="p">,</span> <span class="n">hop_length</span><span class="p">,</span> <span class="n">win_length</span><span class="p">,</span> <span class="n">window</span><span class="p">,</span>  <span class="c1"># type: ignore[attr-defined]</span>
                    <span class="n">normalized</span><span class="p">,</span> <span class="n">onesided</span><span class="p">,</span> <span class="n">return_complex</span><span class="p">)</span></div>


<span class="n">istft</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">istft</span><span class="p">,</span>
    <span class="s2">&quot;istft(input, n_fft, hop_length=None, win_length=None, window=None, center=True, &quot;</span>
    <span class="s2">&quot;normalized=False, onesided=None, length=None, return_complex=False) -&gt; Tensor:</span><span class="se">\n</span><span class="s2">&quot;</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">Inverse short time Fourier Transform. This is expected to be the inverse of :func:`~torch.stft`.</span>

<span class="sd">.. warning::</span>
<span class="sd">    From version 2.1, a warning will be provided if a :attr:`window` is</span>
<span class="sd">    not specified. In a future release, this attribute will be required.</span>
<span class="sd">    Please provide the same window used in the stft call.</span>

<span class="sd">It has the same parameters (+ additional optional parameter of :attr:`length`) and it should return the</span>
<span class="sd">least squares estimation of the original signal. The algorithm will check using the NOLA condition (</span>
<span class="sd">nonzero overlap).</span>

<span class="sd">Important consideration in the parameters :attr:`window` and :attr:`center` so that the envelop</span>
<span class="sd">created by the summation of all the windows is never zero at certain point in time. Specifically,</span>
<span class="sd">:math:`\sum_{t=-\infty}^{\infty} |w|^2[n-t\times hop\_length] \cancel{=} 0`.</span>

<span class="sd">Since :func:`~torch.stft` discards elements at the end of the signal if they do not fit in a frame,</span>
<span class="sd">``istft`` may return a shorter signal than the original signal (can occur if :attr:`center` is False</span>
<span class="sd">since the signal isn&#39;t padded). If `length` is given in the arguments and is longer than expected,</span>
<span class="sd">``istft`` will pad zeros to the end of the returned signal.</span>

<span class="sd">If :attr:`center` is ``True``, then there will be padding e.g. ``&#39;constant&#39;``, ``&#39;reflect&#39;``, etc.</span>
<span class="sd">Left padding can be trimmed off exactly because they can be calculated but right padding cannot be</span>
<span class="sd">calculated without additional information.</span>

<span class="sd">Example: Suppose the last window is:</span>
<span class="sd">``[17, 18, 0, 0, 0]`` vs ``[18, 0, 0, 0, 0]``</span>

<span class="sd">The :attr:`n_fft`, :attr:`hop_length`, :attr:`win_length` are all the same which prevents the calculation</span>
<span class="sd">of right padding. These additional values could be zeros or a reflection of the signal so providing</span>
<span class="sd">:attr:`length` could be useful. If :attr:`length` is ``None`` then padding will be aggressively removed</span>
<span class="sd">(some loss of signal).</span>

<span class="sd">[1] D. W. Griffin and J. S. Lim, &quot;Signal estimation from modified short-time Fourier transform,&quot;</span>
<span class="sd">IEEE Trans. ASSP, vol.32, no.2, pp.236-243, Apr. 1984.</span>

<span class="sd">Args:</span>
<span class="sd">    input (Tensor): The input tensor. Expected to be in the format of :func:`~torch.stft`,</span>
<span class="sd">        output. That is a complex tensor of shape `(B?, N, T)` where</span>

<span class="sd">        - `B?` is an optional batch dimension</span>
<span class="sd">        - `N` is the number of frequency samples, `(n_fft // 2) + 1`</span>
<span class="sd">          for onesided input, or otherwise `n_fft`.</span>
<span class="sd">        - `T` is the number of frames, `1 + length // hop_length` for centered stft,</span>
<span class="sd">          or `1 + (length - n_fft) // hop_length` otherwise.</span>

<span class="sd">        .. versionchanged:: 2.0</span>
<span class="sd">            Real datatype inputs are no longer supported. Input must now have a</span>
<span class="sd">            complex datatype, as returned by ``stft(..., return_complex=True)``.</span>
<span class="sd">    n_fft (int): Size of Fourier transform</span>
<span class="sd">    hop_length (Optional[int]): The distance between neighboring sliding window frames.</span>
<span class="sd">        (Default: ``n_fft // 4``)</span>
<span class="sd">    win_length (Optional[int]): The size of window frame and STFT filter. (Default: ``n_fft``)</span>
<span class="sd">    window (Optional[torch.Tensor]): The optional window function.</span>
<span class="sd">        Shape must be 1d and `&lt;= n_fft`</span>
<span class="sd">        (Default: ``torch.ones(win_length)``)</span>
<span class="sd">    center (bool): Whether :attr:`input` was padded on both sides so that the :math:`t`-th frame is</span>
<span class="sd">        centered at time :math:`t \times \text{hop\_length}`.</span>
<span class="sd">        (Default: ``True``)</span>
<span class="sd">    normalized (bool): Whether the STFT was normalized. (Default: ``False``)</span>
<span class="sd">    onesided (Optional[bool]): Whether the STFT was onesided.</span>
<span class="sd">        (Default: ``True`` if `n_fft != fft_size` in the input size)</span>
<span class="sd">    length (Optional[int]): The amount to trim the signal by (i.e. the</span>
<span class="sd">        original signal length). Defaults to `(T - 1) * hop_length` for</span>
<span class="sd">        centered stft, or `n_fft + (T - 1) * hop_length` otherwise, where `T`</span>
<span class="sd">        is the number of input frames.</span>
<span class="sd">    return_complex (Optional[bool]):</span>
<span class="sd">        Whether the output should be complex, or if the input should be</span>
<span class="sd">        assumed to derive from a real signal and window.</span>
<span class="sd">        Note that this is incompatible with ``onesided=True``.</span>
<span class="sd">        (Default: ``False``)</span>

<span class="sd">Returns:</span>
<span class="sd">    Tensor: Least squares estimation of the original signal of shape `(B?, length)` where</span>
<span class="sd">        `B?` is an optional batch dimension from the input tensor.</span>
<span class="sd">&quot;&quot;&quot;</span><span class="p">)</span>


<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="c1"># These _impl functions return a variable number of tensors as output with</span>
    <span class="c1"># __torch_function__; tuple unpacking is done already rather than being</span>
    <span class="c1"># done by the caller of the _impl function</span>
    <span class="n">_unique_impl_out</span> <span class="o">=</span> <span class="n">Any</span>
<span class="k">else</span><span class="p">:</span>
    <span class="n">_unique_impl_out</span> <span class="o">=</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="k">def</span> <span class="nf">_unique_impl</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="nb">sorted</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">return_inverse</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_counts</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">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="o">-&gt;</span> <span class="n">_unique_impl_out</span><span class="p">:</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;unique(input, sorted=True, return_inverse=False, return_counts=False, dim=None) -&gt; Tuple[Tensor, Tensor, Tensor]</span>

<span class="sd">    Returns the unique elements of the input tensor.</span>

<span class="sd">    .. note:: This function is different from :func:`torch.unique_consecutive` in the sense that</span>
<span class="sd">        this function also eliminates non-consecutive duplicate values.</span>

<span class="sd">    .. note:: Currently in the CUDA implementation and the CPU implementation,</span>
<span class="sd">        `torch.unique` always sort the tensor at the beginning regardless of the `sort` argument.</span>
<span class="sd">        Sorting could be slow, so if your input tensor is already sorted, it is recommended to use</span>
<span class="sd">        :func:`torch.unique_consecutive` which avoids the sorting.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): the input tensor</span>
<span class="sd">        sorted (bool): Whether to sort the unique elements in ascending order</span>
<span class="sd">            before returning as output.</span>
<span class="sd">        return_inverse (bool): Whether to also return the indices for where</span>
<span class="sd">            elements in the original input ended up in the returned unique list.</span>
<span class="sd">        return_counts (bool): Whether to also return the counts for each unique</span>
<span class="sd">            element.</span>
<span class="sd">        dim (int, optional): the dimension to operate upon. If ``None``, the</span>
<span class="sd">            unique of the flattened input is returned. Otherwise, each of the</span>
<span class="sd">            tensors indexed by the given dimension is treated as one of the</span>
<span class="sd">            elements to apply the unique operation upon. See examples for more</span>
<span class="sd">            details. Default: ``None``</span>

<span class="sd">    Returns:</span>
<span class="sd">        (Tensor, Tensor (optional), Tensor (optional)): A tensor or a tuple of tensors containing</span>

<span class="sd">            - **output** (*Tensor*): the output list of unique scalar elements.</span>
<span class="sd">            - **inverse_indices** (*Tensor*): (optional) if</span>
<span class="sd">              :attr:`return_inverse` is True, there will be an additional</span>
<span class="sd">              returned tensor (same shape as input) representing the indices</span>
<span class="sd">              for where elements in the original input map to in the output;</span>
<span class="sd">              otherwise, this function will only return a single tensor.</span>
<span class="sd">            - **counts** (*Tensor*): (optional) if</span>
<span class="sd">              :attr:`return_counts` is True, there will be an additional</span>
<span class="sd">              returned tensor (same shape as output or output.size(dim),</span>
<span class="sd">              if dim was specified) representing the number of occurrences</span>
<span class="sd">              for each unique value or tensor.</span>

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

<span class="sd">        &gt;&gt;&gt; output = torch.unique(torch.tensor([1, 3, 2, 3], dtype=torch.long))</span>
<span class="sd">        &gt;&gt;&gt; output</span>
<span class="sd">        tensor([1, 2, 3])</span>

<span class="sd">        &gt;&gt;&gt; output, inverse_indices = torch.unique(</span>
<span class="sd">        ...     torch.tensor([1, 3, 2, 3], dtype=torch.long), sorted=True, return_inverse=True)</span>
<span class="sd">        &gt;&gt;&gt; output</span>
<span class="sd">        tensor([1, 2, 3])</span>
<span class="sd">        &gt;&gt;&gt; inverse_indices</span>
<span class="sd">        tensor([0, 2, 1, 2])</span>

<span class="sd">        &gt;&gt;&gt; output, inverse_indices = torch.unique(</span>
<span class="sd">        ...     torch.tensor([[1, 3], [2, 3]], dtype=torch.long), sorted=True, return_inverse=True)</span>
<span class="sd">        &gt;&gt;&gt; output</span>
<span class="sd">        tensor([1, 2, 3])</span>
<span class="sd">        &gt;&gt;&gt; inverse_indices</span>
<span class="sd">        tensor([[0, 2],</span>
<span class="sd">                [1, 2]])</span>

<span class="sd">        &gt;&gt;&gt; a = torch.tensor([</span>
<span class="sd">        ...     [</span>
<span class="sd">        ...         [1, 1, 0, 0],</span>
<span class="sd">        ...         [1, 1, 0, 0],</span>
<span class="sd">        ...         [0, 0, 1, 1],</span>
<span class="sd">        ...     ],</span>
<span class="sd">        ...     [</span>
<span class="sd">        ...         [0, 0, 1, 1],</span>
<span class="sd">        ...         [0, 0, 1, 1],</span>
<span class="sd">        ...         [1, 1, 1, 1],</span>
<span class="sd">        ...     ],</span>
<span class="sd">        ...     [</span>
<span class="sd">        ...         [1, 1, 0, 0],</span>
<span class="sd">        ...         [1, 1, 0, 0],</span>
<span class="sd">        ...         [0, 0, 1, 1],</span>
<span class="sd">        ...     ],</span>
<span class="sd">        ... ])</span>

<span class="sd">        &gt;&gt;&gt; # If we call `torch.unique(a, dim=0)`, each of the tensors `a[idx, :, :]`</span>
<span class="sd">        &gt;&gt;&gt; # will be compared. We can see that `a[0, :, :]` and `a[2, :, :]` match</span>
<span class="sd">        &gt;&gt;&gt; # each other, so one of them will be removed.</span>
<span class="sd">        &gt;&gt;&gt; (a[0, :, :] == a[2, :, :]).all()</span>
<span class="sd">        tensor(True)</span>
<span class="sd">        &gt;&gt;&gt; a_unique_dim0 = torch.unique(a, dim=0)</span>
<span class="sd">        &gt;&gt;&gt; a_unique_dim0</span>
<span class="sd">        tensor([[[0, 0, 1, 1],</span>
<span class="sd">                 [0, 0, 1, 1],</span>
<span class="sd">                 [1, 1, 1, 1]],</span>
<span class="sd">                [[1, 1, 0, 0],</span>
<span class="sd">                 [1, 1, 0, 0],</span>
<span class="sd">                 [0, 0, 1, 1]]])</span>

<span class="sd">        &gt;&gt;&gt; # Notice which sub-tensors from `a` match with the sub-tensors from</span>
<span class="sd">        &gt;&gt;&gt; # `a_unique_dim0`:</span>
<span class="sd">        &gt;&gt;&gt; (a_unique_dim0[0, :, :] == a[1, :, :]).all()</span>
<span class="sd">        tensor(True)</span>
<span class="sd">        &gt;&gt;&gt; (a_unique_dim0[1, :, :] == a[0, :, :]).all()</span>
<span class="sd">        tensor(True)</span>

<span class="sd">        &gt;&gt;&gt; # For `torch.unique(a, dim=1)`, each of the tensors `a[:, idx, :]` are</span>
<span class="sd">        &gt;&gt;&gt; # compared. `a[:, 0, :]` and `a[:, 1, :]` match each other, so one of</span>
<span class="sd">        &gt;&gt;&gt; # them will be removed.</span>
<span class="sd">        &gt;&gt;&gt; (a[:, 0, :] == a[:, 1, :]).all()</span>
<span class="sd">        tensor(True)</span>
<span class="sd">        &gt;&gt;&gt; torch.unique(a, dim=1)</span>
<span class="sd">        tensor([[[0, 0, 1, 1],</span>
<span class="sd">                 [1, 1, 0, 0]],</span>
<span class="sd">                [[1, 1, 1, 1],</span>
<span class="sd">                 [0, 0, 1, 1]],</span>
<span class="sd">                [[0, 0, 1, 1],</span>
<span class="sd">                 [1, 1, 0, 0]]])</span>

<span class="sd">        &gt;&gt;&gt; # For `torch.unique(a, dim=2)`, the tensors `a[:, :, idx]` are compared.</span>
<span class="sd">        &gt;&gt;&gt; # `a[:, :, 0]` and `a[:, :, 1]` match each other. Also, `a[:, :, 2]` and</span>
<span class="sd">        &gt;&gt;&gt; # `a[:, :, 3]` match each other as well. So in this case, two of the</span>
<span class="sd">        &gt;&gt;&gt; # sub-tensors will be removed.</span>
<span class="sd">        &gt;&gt;&gt; (a[:, :, 0] == a[:, :, 1]).all()</span>
<span class="sd">        tensor(True)</span>
<span class="sd">        &gt;&gt;&gt; (a[:, :, 2] == a[:, :, 3]).all()</span>
<span class="sd">        tensor(True)</span>
<span class="sd">        &gt;&gt;&gt; torch.unique(a, dim=2)</span>
<span class="sd">        tensor([[[0, 1],</span>
<span class="sd">                 [0, 1],</span>
<span class="sd">                 [1, 0]],</span>
<span class="sd">                [[1, 0],</span>
<span class="sd">                 [1, 0],</span>
<span class="sd">                 [1, 1]],</span>
<span class="sd">                [[0, 1],</span>
<span class="sd">                 [0, 1],</span>
<span class="sd">                 [1, 0]]])</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">unique</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="nb">sorted</span><span class="o">=</span><span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="n">return_inverse</span><span class="p">,</span>
            <span class="n">return_counts</span><span class="o">=</span><span class="n">return_counts</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">dim</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">counts</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">unique_dim</span><span class="p">(</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="n">dim</span><span class="p">,</span>
            <span class="nb">sorted</span><span class="o">=</span><span class="nb">sorted</span><span class="p">,</span>
            <span class="n">return_inverse</span><span class="o">=</span><span class="n">return_inverse</span><span class="p">,</span>
            <span class="n">return_counts</span><span class="o">=</span><span class="n">return_counts</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">counts</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_unique2</span><span class="p">(</span>
            <span class="nb">input</span><span class="p">,</span>
            <span class="nb">sorted</span><span class="o">=</span><span class="nb">sorted</span><span class="p">,</span>
            <span class="n">return_inverse</span><span class="o">=</span><span class="n">return_inverse</span><span class="p">,</span>
            <span class="n">return_counts</span><span class="o">=</span><span class="n">return_counts</span><span class="p">,</span>
        <span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">counts</span>


<span class="k">def</span> <span class="nf">_unique_consecutive_impl</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">return_inverse</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_counts</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">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="o">-&gt;</span> <span class="n">_unique_impl_out</span><span class="p">:</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Eliminates all but the first element from every consecutive group of equivalent elements.</span>

<span class="sd">    .. note:: This function is different from :func:`torch.unique` in the sense that this function</span>
<span class="sd">        only eliminates consecutive duplicate values. This semantics is similar to `std::unique`</span>
<span class="sd">        in C++.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): the input tensor</span>
<span class="sd">        return_inverse (bool): Whether to also return the indices for where</span>
<span class="sd">            elements in the original input ended up in the returned unique list.</span>
<span class="sd">        return_counts (bool): Whether to also return the counts for each unique</span>
<span class="sd">            element.</span>
<span class="sd">        dim (int): the dimension to apply unique. If ``None``, the unique of the</span>
<span class="sd">            flattened input is returned. default: ``None``</span>

<span class="sd">    Returns:</span>
<span class="sd">        (Tensor, Tensor (optional), Tensor (optional)): A tensor or a tuple of tensors containing</span>

<span class="sd">            - **output** (*Tensor*): the output list of unique scalar elements.</span>
<span class="sd">            - **inverse_indices** (*Tensor*): (optional) if</span>
<span class="sd">              :attr:`return_inverse` is True, there will be an additional</span>
<span class="sd">              returned tensor (same shape as input) representing the indices</span>
<span class="sd">              for where elements in the original input map to in the output;</span>
<span class="sd">              otherwise, this function will only return a single tensor.</span>
<span class="sd">            - **counts** (*Tensor*): (optional) if</span>
<span class="sd">              :attr:`return_counts` is True, there will be an additional</span>
<span class="sd">              returned tensor (same shape as output or output.size(dim),</span>
<span class="sd">              if dim was specified) representing the number of occurrences</span>
<span class="sd">              for each unique value or tensor.</span>

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

<span class="sd">        &gt;&gt;&gt; x = torch.tensor([1, 1, 2, 2, 3, 1, 1, 2])</span>
<span class="sd">        &gt;&gt;&gt; output = torch.unique_consecutive(x)</span>
<span class="sd">        &gt;&gt;&gt; output</span>
<span class="sd">        tensor([1, 2, 3, 1, 2])</span>

<span class="sd">        &gt;&gt;&gt; output, inverse_indices = torch.unique_consecutive(x, return_inverse=True)</span>
<span class="sd">        &gt;&gt;&gt; output</span>
<span class="sd">        tensor([1, 2, 3, 1, 2])</span>
<span class="sd">        &gt;&gt;&gt; inverse_indices</span>
<span class="sd">        tensor([0, 0, 1, 1, 2, 3, 3, 4])</span>

<span class="sd">        &gt;&gt;&gt; output, counts = torch.unique_consecutive(x, return_counts=True)</span>
<span class="sd">        &gt;&gt;&gt; output</span>
<span class="sd">        tensor([1, 2, 3, 1, 2])</span>
<span class="sd">        &gt;&gt;&gt; counts</span>
<span class="sd">        tensor([2, 2, 1, 2, 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">unique_consecutive</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">return_inverse</span><span class="o">=</span><span class="n">return_inverse</span><span class="p">,</span>
            <span class="n">return_counts</span><span class="o">=</span><span class="n">return_counts</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">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">counts</span> <span class="o">=</span> <span class="n">_VF</span><span class="o">.</span><span class="n">unique_consecutive</span><span class="p">(</span>  <span class="c1"># type: ignore[attr-defined]</span>
        <span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="n">return_counts</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">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">counts</span>


<span class="k">def</span> <span class="nf">_return_counts</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, bool, Optional[int]) -&gt; Tuple[Tensor, Tensor]</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">_unique_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>

    <span class="n">output</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">counts</span> <span class="o">=</span> <span class="n">_unique_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">counts</span>


<span class="k">def</span> <span class="nf">_return_output</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, bool, Optional[int]) -&gt; Tensor</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">_unique_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>

    <span class="n">output</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">_unique_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span>


<span class="k">def</span> <span class="nf">_return_inverse</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, bool, Optional[int]) -&gt; Tuple[Tensor, Tensor]</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">_unique_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>

    <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">_unique_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="nb">sorted</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span>


<span class="n">_return_inverse_false</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="s1">&#39;return_counts&#39;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">3</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">_return_counts</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_return_output</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="s1">&#39;unique&#39;</span><span class="p">)</span>

<span class="n">_return_inverse_true</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="s1">&#39;return_counts&#39;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">3</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">_unique_impl</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_return_inverse</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="s1">&#39;unique&#39;</span><span class="p">)</span>

<span class="c1"># The return type of unique depends on `return_inverse`, and `return_counts` so in order to</span>
<span class="c1"># resolve the output type in TorchScript we need to statically know the value of both parameters</span>

<span class="n">unique</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="s1">&#39;return_inverse&#39;</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">_return_inverse_true</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_return_inverse_false</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="s1">&#39;unique&#39;</span><span class="p">)</span>
<span class="n">unique</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">_unique_impl</span><span class="o">.</span><span class="vm">__doc__</span>


<span class="k">def</span> <span class="nf">_consecutive_return_counts</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, Optional[int]) -&gt; Tuple[Tensor, Tensor]</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">_unique_consecutive_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>

    <span class="n">output</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">counts</span> <span class="o">=</span> <span class="n">_unique_consecutive_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">counts</span>


<span class="k">def</span> <span class="nf">_consecutive_return_output</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, Optional[int]) -&gt; Tensor</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">_unique_consecutive_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>

    <span class="n">output</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">_unique_consecutive_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span>


<span class="k">def</span> <span class="nf">_consecutive_return_inverse</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">return_counts</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, Optional[int]) -&gt; Tuple[Tensor, Tensor]</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">_unique_consecutive_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>

    <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">_unique_consecutive_impl</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">return_inverse</span><span class="p">,</span> <span class="n">return_counts</span><span class="p">,</span> <span class="n">dim</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">inverse_indices</span>


<span class="n">_consecutive_return_inverse_false</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="s1">&#39;return_counts&#39;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">1</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">_consecutive_return_counts</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_consecutive_return_output</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="s1">&#39;unique_consecutive&#39;</span><span class="p">)</span>

<span class="n">_consecutive_return_inverse_true</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="s1">&#39;return_counts&#39;</span><span class="p">,</span>
    <span class="n">arg_index</span><span class="o">=</span><span class="mi">1</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">_unique_consecutive_impl</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_consecutive_return_inverse</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="s1">&#39;unique_consecutive&#39;</span><span class="p">)</span>

<span class="c1"># The return type of unique depends on `return_inverse`, and `return_counts` so in order to</span>
<span class="c1"># resolve the output type in TorchScript we need to statically know the value of both parameters</span>

<span class="n">unique_consecutive</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="s1">&#39;return_inverse&#39;</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">_consecutive_return_inverse_true</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_consecutive_return_inverse_false</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="s1">&#39;unique_consecutive&#39;</span><span class="p">)</span>
<span class="n">unique_consecutive</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">_unique_consecutive_impl</span><span class="o">.</span><span class="vm">__doc__</span>

<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="k">pass</span>
    <span class="c1"># There&#39;s no good way to use this type annotation without breaking JIT</span>
    <span class="c1"># overloads. So leave untyped for mypy for now.</span>
<span class="k">else</span><span class="p">:</span>
    <span class="nd">@overload</span>
    <span class="k">def</span> <span class="nf">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">2</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">torch</span><span class="o">.</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="k">pass</span>

    <span class="nd">@overload</span>  <span class="c1"># noqa: F811</span>
    <span class="k">def</span> <span class="nf">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims</span><span class="p">:</span> <span class="n">Tuple</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">List</span><span class="p">[</span><span class="nb">int</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">torch</span><span class="o">.</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="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">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims</span><span class="p">:</span> <span class="n">List</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">out</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</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">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="n">out</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
        <span class="k">pass</span>


<div class="viewcode-block" id="tensordot"><a class="viewcode-back" href="../../generated/torch.tensordot.html#torch.tensordot">[docs]</a><span class="k">def</span> <span class="nf">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims</span><span class="o">=</span><span class="mi">2</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">torch</span><span class="o">.</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="c1"># noqa: F811</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Returns a contraction of a and b over multiple dimensions.</span>

<span class="sd">    :attr:`tensordot` implements a generalized matrix product.</span>

<span class="sd">    Args:</span>
<span class="sd">      a (Tensor): Left tensor to contract</span>
<span class="sd">      b (Tensor): Right tensor to contract</span>
<span class="sd">      dims (int or Tuple[List[int], List[int]] or List[List[int]] containing two lists or Tensor): number of dimensions to</span>
<span class="sd">         contract or explicit lists of dimensions for :attr:`a` and</span>
<span class="sd">         :attr:`b` respectively</span>

<span class="sd">    When called with a non-negative integer argument :attr:`dims` = :math:`d`, and</span>
<span class="sd">    the number of dimensions of :attr:`a` and :attr:`b` is :math:`m` and :math:`n`,</span>
<span class="sd">    respectively, :func:`~torch.tensordot` computes</span>

<span class="sd">    .. math::</span>
<span class="sd">        r_{i_0,...,i_{m-d}, i_d,...,i_n}</span>
<span class="sd">          = \sum_{k_0,...,k_{d-1}} a_{i_0,...,i_{m-d},k_0,...,k_{d-1}} \times b_{k_0,...,k_{d-1}, i_d,...,i_n}.</span>

<span class="sd">    When called with :attr:`dims` of the list form, the given dimensions will be contracted</span>
<span class="sd">    in place of the last :math:`d` of :attr:`a` and the first :math:`d` of :math:`b`. The sizes</span>
<span class="sd">    in these dimensions must match, but :func:`~torch.tensordot` will deal with broadcasted</span>
<span class="sd">    dimensions.</span>

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

<span class="sd">        &gt;&gt;&gt; a = torch.arange(60.).reshape(3, 4, 5)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.arange(24.).reshape(4, 3, 2)</span>
<span class="sd">        &gt;&gt;&gt; torch.tensordot(a, b, dims=([1, 0], [0, 1]))</span>
<span class="sd">        tensor([[4400., 4730.],</span>
<span class="sd">                [4532., 4874.],</span>
<span class="sd">                [4664., 5018.],</span>
<span class="sd">                [4796., 5162.],</span>
<span class="sd">                [4928., 5306.]])</span>

<span class="sd">        &gt;&gt;&gt; # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)</span>
<span class="sd">        &gt;&gt;&gt; a = torch.randn(3, 4, 5, device=&#39;cuda&#39;)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.randn(4, 5, 6, device=&#39;cuda&#39;)</span>
<span class="sd">        &gt;&gt;&gt; c = torch.tensordot(a, b, dims=2).cpu()</span>
<span class="sd">        tensor([[ 8.3504, -2.5436,  6.2922,  2.7556, -1.0732,  3.2741],</span>
<span class="sd">                [ 3.3161,  0.0704,  5.0187, -0.4079, -4.3126,  4.8744],</span>
<span class="sd">                [ 0.8223,  3.9445,  3.2168, -0.2400,  3.4117,  1.7780]])</span>

<span class="sd">        &gt;&gt;&gt; a = torch.randn(3, 5, 4, 6)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.randn(6, 4, 5, 3)</span>
<span class="sd">        &gt;&gt;&gt; torch.tensordot(a, b, dims=([2, 1, 3], [1, 2, 0]))</span>
<span class="sd">        tensor([[  7.7193,  -2.4867, -10.3204],</span>
<span class="sd">                [  1.5513, -14.4737,  -6.5113],</span>
<span class="sd">                [ -0.2850,   4.2573,  -3.5997]])</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">a</span><span class="p">,</span> <span class="n">b</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">tensordot</span><span class="p">,</span> <span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">),</span> <span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims</span><span class="o">=</span><span class="n">dims</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="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">,</span> <span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;tensordot expects dims to be int or &quot;</span>
                           <span class="o">+</span> <span class="s2">&quot;Tuple[List[int], List[int]] or &quot;</span>
                           <span class="o">+</span> <span class="s2">&quot;List[List[int]] containing two lists, but got &quot;</span>
                           <span class="o">+</span> <span class="sa">f</span><span class="s2">&quot;dims=</span><span class="si">{</span><span class="n">dims</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

    <span class="n">dims_a</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="p">[]</span>
    <span class="n">dims_b</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="p">[]</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
        <span class="n">dims_a</span><span class="p">,</span> <span class="n">dims_b</span> <span class="o">=</span> <span class="n">dims</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="n">num_elements</span> <span class="o">=</span> <span class="n">dims</span><span class="o">.</span><span class="n">numel</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">num_elements</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">assert</span> <span class="n">dims</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">2</span>
            <span class="n">dims_a</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="n">dims</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span>
            <span class="n">dims_b</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="n">dims</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">dims_val</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">dims</span><span class="o">.</span><span class="n">item</span><span class="p">())</span>
            <span class="k">if</span> <span class="n">dims_val</span> <span class="o">&lt;</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="sa">f</span><span class="s2">&quot;tensordot expects dims &gt;= 0, but got dims=</span><span class="si">{</span><span class="n">dims</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
            <span class="n">dims_a</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="o">-</span><span class="n">dims_val</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
            <span class="n">dims_b</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">dims_val</span><span class="p">))</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dims</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
        <span class="k">if</span> <span class="n">dims</span> <span class="o">&lt;</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="sa">f</span><span class="s2">&quot;tensordot expects dims &gt;= 0, but got dims=</span><span class="si">{</span><span class="n">dims</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">dims</span> <span class="o">&gt;</span> <span class="nb">min</span><span class="p">(</span><span class="n">a</span><span class="o">.</span><span class="n">dim</span><span class="p">(),</span> <span class="n">b</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="sa">f</span><span class="s2">&quot;tensordot expects dims &lt; ndim_a or ndim_b, but got dims=</span><span class="si">{</span><span class="n">dims</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
        <span class="n">dims_a</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="o">-</span><span class="n">dims</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
        <span class="n">dims_b</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">dims</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="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims_a</span><span class="p">,</span> <span class="n">dims_b</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">tensordot</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">,</span> <span class="n">dims_a</span><span class="p">,</span> <span class="n">dims_b</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="c1"># type: ignore[attr-defined]</span></div>


<div class="viewcode-block" id="cartesian_prod"><a class="viewcode-back" href="../../generated/torch.cartesian_prod.html#torch.cartesian_prod">[docs]</a><span class="k">def</span> <span class="nf">cartesian_prod</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">:</span> <span class="n">Tensor</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tensor</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Do cartesian product of the given sequence of tensors. The behavior is similar to</span>
<span class="sd">    python&#39;s `itertools.product`.</span>

<span class="sd">    Args:</span>
<span class="sd">        *tensors: any number of 1 dimensional tensors.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Tensor: A tensor equivalent to converting all the input tensors into lists,</span>
<span class="sd">        do `itertools.product` on these lists, and finally convert the resulting list</span>
<span class="sd">        into tensor.</span>

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

<span class="sd">        &gt;&gt;&gt; import itertools</span>
<span class="sd">        &gt;&gt;&gt; a = [1, 2, 3]</span>
<span class="sd">        &gt;&gt;&gt; b = [4, 5]</span>
<span class="sd">        &gt;&gt;&gt; list(itertools.product(a, b))</span>
<span class="sd">        [(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]</span>
<span class="sd">        &gt;&gt;&gt; tensor_a = torch.tensor(a)</span>
<span class="sd">        &gt;&gt;&gt; tensor_b = torch.tensor(b)</span>
<span class="sd">        &gt;&gt;&gt; torch.cartesian_prod(tensor_a, tensor_b)</span>
<span class="sd">        tensor([[1, 4],</span>
<span class="sd">                [1, 5],</span>
<span class="sd">                [2, 4],</span>
<span class="sd">                [2, 5],</span>
<span class="sd">                [3, 4],</span>
<span class="sd">                [3, 5]])</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">cartesian_prod</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">cartesian_prod</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<div class="viewcode-block" id="block_diag"><a class="viewcode-back" href="../../generated/torch.block_diag.html#torch.block_diag">[docs]</a><span class="k">def</span> <span class="nf">block_diag</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Create a block diagonal matrix from provided tensors.</span>

<span class="sd">    Args:</span>
<span class="sd">        *tensors: One or more tensors with 0, 1, or 2 dimensions.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Tensor: A 2 dimensional tensor with all the input tensors arranged in</span>
<span class="sd">        order such that their upper left and lower right corners are</span>
<span class="sd">        diagonally adjacent. All other elements are set to 0.</span>

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

<span class="sd">        &gt;&gt;&gt; import torch</span>
<span class="sd">        &gt;&gt;&gt; A = torch.tensor([[0, 1], [1, 0]])</span>
<span class="sd">        &gt;&gt;&gt; B = torch.tensor([[3, 4, 5], [6, 7, 8]])</span>
<span class="sd">        &gt;&gt;&gt; C = torch.tensor(7)</span>
<span class="sd">        &gt;&gt;&gt; D = torch.tensor([1, 2, 3])</span>
<span class="sd">        &gt;&gt;&gt; E = torch.tensor([[4], [5], [6]])</span>
<span class="sd">        &gt;&gt;&gt; torch.block_diag(A, B, C, D, E)</span>
<span class="sd">        tensor([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0],</span>
<span class="sd">                [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],</span>
<span class="sd">                [0, 0, 3, 4, 5, 0, 0, 0, 0, 0],</span>
<span class="sd">                [0, 0, 6, 7, 8, 0, 0, 0, 0, 0],</span>
<span class="sd">                [0, 0, 0, 0, 0, 7, 0, 0, 0, 0],</span>
<span class="sd">                [0, 0, 0, 0, 0, 0, 1, 2, 3, 0],</span>
<span class="sd">                [0, 0, 0, 0, 0, 0, 0, 0, 0, 4],</span>
<span class="sd">                [0, 0, 0, 0, 0, 0, 0, 0, 0, 5],</span>
<span class="sd">                [0, 0, 0, 0, 0, 0, 0, 0, 0, 6]])</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">block_diag</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</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">_VariableFunctions</span><span class="o">.</span><span class="n">block_diag</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<div class="viewcode-block" id="cdist"><a class="viewcode-back" href="../../generated/torch.cdist.html#torch.cdist">[docs]</a><span class="k">def</span> <span class="nf">cdist</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="mf">2.</span><span class="p">,</span> <span class="n">compute_mode</span><span class="o">=</span><span class="s1">&#39;use_mm_for_euclid_dist_if_necessary&#39;</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, Tensor, float, str) -&gt; (Tensor)</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Computes batched the p-norm distance between each pair of the two collections of row vectors.</span>

<span class="sd">    Args:</span>
<span class="sd">        x1 (Tensor): input tensor of shape :math:`B \times P \times M`.</span>
<span class="sd">        x2 (Tensor): input tensor of shape :math:`B \times R \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">        compute_mode:</span>
<span class="sd">            &#39;use_mm_for_euclid_dist_if_necessary&#39; - will use matrix multiplication approach to calculate</span>
<span class="sd">            euclidean distance (p = 2) if P &gt; 25 or R &gt; 25</span>
<span class="sd">            &#39;use_mm_for_euclid_dist&#39; - will always use matrix multiplication approach to calculate</span>
<span class="sd">            euclidean distance (p = 2)</span>
<span class="sd">            &#39;donot_use_mm_for_euclid_dist&#39; - will never use matrix multiplication approach to calculate</span>
<span class="sd">            euclidean distance (p = 2)</span>
<span class="sd">            Default: use_mm_for_euclid_dist_if_necessary.</span>

<span class="sd">    If x1 has shape :math:`B \times P \times M` and x2 has shape :math:`B \times R \times M` then the</span>
<span class="sd">    output will have shape :math:`B \times P \times R`.</span>

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

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

<span class="sd">        &gt;&gt;&gt; a = torch.tensor([[0.9041,  0.0196], [-0.3108, -2.4423], [-0.4821,  1.059]])</span>
<span class="sd">        &gt;&gt;&gt; a</span>
<span class="sd">        tensor([[ 0.9041,  0.0196],</span>
<span class="sd">                [-0.3108, -2.4423],</span>
<span class="sd">                [-0.4821,  1.0590]])</span>
<span class="sd">        &gt;&gt;&gt; b = torch.tensor([[-2.1763, -0.4713], [-0.6986,  1.3702]])</span>
<span class="sd">        &gt;&gt;&gt; b</span>
<span class="sd">        tensor([[-2.1763, -0.4713],</span>
<span class="sd">                [-0.6986,  1.3702]])</span>
<span class="sd">        &gt;&gt;&gt; torch.cdist(a, b, p=2)</span>
<span class="sd">        tensor([[3.1193, 2.0959],</span>
<span class="sd">                [2.7138, 3.8322],</span>
<span class="sd">                [2.2830, 0.3791]])</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">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">cdist</span><span class="p">,</span> <span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">),</span> <span class="n">x1</span><span class="p">,</span> <span class="n">x2</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">compute_mode</span><span class="o">=</span><span class="n">compute_mode</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">compute_mode</span> <span class="o">==</span> <span class="s1">&#39;use_mm_for_euclid_dist_if_necessary&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">cdist</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>
    <span class="k">elif</span> <span class="n">compute_mode</span> <span class="o">==</span> <span class="s1">&#39;use_mm_for_euclid_dist&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">cdist</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>
    <span class="k">elif</span> <span class="n">compute_mode</span> <span class="o">==</span> <span class="s1">&#39;donot_use_mm_for_euclid_dist&#39;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">cdist</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">,</span> <span class="n">p</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</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="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">compute_mode</span><span class="si">}</span><span class="s2"> is not a valid value for compute_mode&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="atleast_1d"><a class="viewcode-back" href="../../generated/torch.atleast_1d.html#torch.atleast_1d">[docs]</a><span class="k">def</span> <span class="nf">atleast_1d</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a 1-dimensional view of each input tensor with zero dimensions.</span>
<span class="sd">    Input tensors with one or more dimensions are returned as-is.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor or list of Tensors)</span>

<span class="sd">    Returns:</span>
<span class="sd">        output (Tensor or tuple of Tensors)</span>

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

<span class="sd">        &gt;&gt;&gt; x = torch.arange(2)</span>
<span class="sd">        &gt;&gt;&gt; x</span>
<span class="sd">        tensor([0, 1])</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_1d(x)</span>
<span class="sd">        tensor([0, 1])</span>
<span class="sd">        &gt;&gt;&gt; x = torch.tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; x</span>
<span class="sd">        tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_1d(x)</span>
<span class="sd">        tensor([1.])</span>
<span class="sd">        &gt;&gt;&gt; x = torch.tensor(0.5)</span>
<span class="sd">        &gt;&gt;&gt; y = torch.tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_1d((x, y))</span>
<span class="sd">        (tensor([0.5000]), tensor([1.]))</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">atleast_1d</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">atleast_1d</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<div class="viewcode-block" id="atleast_2d"><a class="viewcode-back" href="../../generated/torch.atleast_2d.html#torch.atleast_2d">[docs]</a><span class="k">def</span> <span class="nf">atleast_2d</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a 2-dimensional view of each input tensor with zero dimensions.</span>
<span class="sd">    Input tensors with two or more dimensions are returned as-is.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor or list of Tensors)</span>

<span class="sd">    Returns:</span>
<span class="sd">        output (Tensor or tuple of Tensors)</span>

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

<span class="sd">        &gt;&gt;&gt; x = torch.tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; x</span>
<span class="sd">        tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_2d(x)</span>
<span class="sd">        tensor([[1.]])</span>
<span class="sd">        &gt;&gt;&gt; x = torch.arange(4).view(2, 2)</span>
<span class="sd">        &gt;&gt;&gt; x</span>
<span class="sd">        tensor([[0, 1],</span>
<span class="sd">                [2, 3]])</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_2d(x)</span>
<span class="sd">        tensor([[0, 1],</span>
<span class="sd">                [2, 3]])</span>
<span class="sd">        &gt;&gt;&gt; x = torch.tensor(0.5)</span>
<span class="sd">        &gt;&gt;&gt; y = torch.tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_2d((x, y))</span>
<span class="sd">        (tensor([[0.5000]]), tensor([[1.]]))</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">atleast_2d</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">atleast_2d</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<div class="viewcode-block" id="atleast_3d"><a class="viewcode-back" href="../../generated/torch.atleast_3d.html#torch.atleast_3d">[docs]</a><span class="k">def</span> <span class="nf">atleast_3d</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns a 3-dimensional view of each input tensor with zero dimensions.</span>
<span class="sd">    Input tensors with three or more dimensions are returned as-is.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor or list of Tensors)</span>

<span class="sd">    Returns:</span>
<span class="sd">        output (Tensor or tuple of Tensors)</span>

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

<span class="sd">        &gt;&gt;&gt; x = torch.tensor(0.5)</span>
<span class="sd">        &gt;&gt;&gt; x</span>
<span class="sd">        tensor(0.5000)</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_3d(x)</span>
<span class="sd">        tensor([[[0.5000]]])</span>
<span class="sd">        &gt;&gt;&gt; y = torch.arange(4).view(2, 2)</span>
<span class="sd">        &gt;&gt;&gt; y</span>
<span class="sd">        tensor([[0, 1],</span>
<span class="sd">                [2, 3]])</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_3d(y)</span>
<span class="sd">        tensor([[[0],</span>
<span class="sd">                 [1]],</span>
<span class="sd">                &lt;BLANKLINE&gt;</span>
<span class="sd">                [[2],</span>
<span class="sd">                 [3]]])</span>
<span class="sd">        &gt;&gt;&gt; x = torch.tensor(1).view(1, 1, 1)</span>
<span class="sd">        &gt;&gt;&gt; x</span>
<span class="sd">        tensor([[[1]]])</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_3d(x)</span>
<span class="sd">        tensor([[[1]]])</span>
<span class="sd">        &gt;&gt;&gt; x = torch.tensor(0.5)</span>
<span class="sd">        &gt;&gt;&gt; y = torch.tensor(1.)</span>
<span class="sd">        &gt;&gt;&gt; torch.atleast_3d((x, y))</span>
<span class="sd">        (tensor([[[0.5000]]]), tensor([[[1.]]]))</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">tensors</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">atleast_3d</span><span class="p">,</span> <span class="n">tensors</span><span class="p">,</span> <span class="o">*</span><span class="n">tensors</span><span class="p">)</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">tensors</span> <span class="o">=</span> <span class="n">tensors</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">atleast_3d</span><span class="p">(</span><span class="n">tensors</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>


<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="k">pass</span>
    <span class="c1"># There&#39;s no good way to use this type annotation; cannot rename norm() to</span>
    <span class="c1"># _norm_impl() in a way that doesn&#39;t break JIT overloads. So leave untyped</span>
    <span class="c1"># for mypy for now.</span>
    <span class="c1">#    def norm(input: Tensor,</span>
    <span class="c1">#             p: Optional[Union[str, Number]] = &quot;fro&quot;,</span>
    <span class="c1">#             dim: Optional[Union[int, List[int]]] = None,</span>
    <span class="c1">#             keepdim: bool = False,</span>
    <span class="c1">#             out: Optional[Tensor] = None,</span>
    <span class="c1">#             dtype: _dtype = None) -&gt; Tensor:</span>
    <span class="c1">#        return _norm_impl(input, p, dim, keepdim, out, dtype)</span>
<span class="k">else</span><span class="p">:</span>
    <span class="c1"># TODO: type dim as BroadcastingList when</span>
    <span class="c1"># https://github.com/pytorch/pytorch/issues/33782 is fixed</span>
    <span class="nd">@overload</span>
    <span class="k">def</span> <span class="nf">norm</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="s2">&quot;fro&quot;</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="c1"># type: (Tensor, str, Optional[List[int]], bool, Optional[Tensor], Optional[int]) -&gt; Tensor</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">norm</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="s2">&quot;fro&quot;</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
        <span class="c1"># type: (Tensor, Optional[number], Optional[List[int]], bool, Optional[Tensor], Optional[int]) -&gt; Tensor</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">norm</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="s2">&quot;fro&quot;</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
        <span class="c1"># type: (Tensor, Optional[number], Optional[int], bool, Optional[Tensor], Optional[int]) -&gt; Tensor</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">norm</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="s2">&quot;fro&quot;</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
        <span class="c1"># type: (Tensor, str, Optional[int], bool, Optional[Tensor], Optional[int]) -&gt; Tensor</span>
        <span class="k">pass</span>


<div class="viewcode-block" id="norm"><a class="viewcode-back" href="../../generated/torch.norm.html#torch.norm">[docs]</a><span class="k">def</span> <span class="nf">norm</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">Optional</span><span class="p">[</span><span class="n">Union</span><span class="p">[</span><span class="nb">float</span><span class="p">,</span> <span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="s2">&quot;fro&quot;</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>  <span class="c1"># noqa: F811</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Returns the matrix norm or vector norm of a given tensor.</span>

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

<span class="sd">        torch.norm is deprecated and may be removed in a future PyTorch release.</span>
<span class="sd">        Its documentation and behavior may be incorrect, and it is no longer</span>
<span class="sd">        actively maintained.</span>

<span class="sd">        Use :func:`torch.linalg.vector_norm` when computing vector norms and</span>
<span class="sd">        :func:`torch.linalg.matrix_norm` when computing matrix norms.</span>
<span class="sd">        For a function with a similar behavior as this one see :func:`torch.linalg.norm`.</span>
<span class="sd">        Note, however, the signature for these functions is slightly different than the</span>
<span class="sd">        signature for ``torch.norm``.</span>

<span class="sd">    Args:</span>
<span class="sd">        input (Tensor): The input tensor. Its data type must be either a floating</span>
<span class="sd">            point or complex type. For complex inputs, the norm is calculated using the</span>
<span class="sd">            absolute value of each element. If the input is complex and neither</span>
<span class="sd">            :attr:`dtype` nor :attr:`out` is specified, the result&#39;s data type will</span>
<span class="sd">            be the corresponding floating point type (e.g. float if :attr:`input` is</span>
<span class="sd">            complexfloat).</span>

<span class="sd">        p (int, float, inf, -inf, &#39;fro&#39;, &#39;nuc&#39;, optional): the order of norm. Default: ``&#39;fro&#39;``</span>
<span class="sd">            The following norms can be calculated:</span>

<span class="sd">            ======  ==============  ==========================</span>
<span class="sd">            ord     matrix norm     vector norm</span>
<span class="sd">            ======  ==============  ==========================</span>
<span class="sd">            &#39;fro&#39;   Frobenius norm  --</span>
<span class="sd">            &#39;nuc&#39;   nuclear norm    --</span>
<span class="sd">            Number  --              sum(abs(x)**ord)**(1./ord)</span>
<span class="sd">            ======  ==============  ==========================</span>

<span class="sd">            The vector norm can be calculated across any number of dimensions.</span>
<span class="sd">            The corresponding dimensions of :attr:`input` are flattened into</span>
<span class="sd">            one dimension, and the norm is calculated on the flattened</span>
<span class="sd">            dimension.</span>

<span class="sd">            Frobenius norm produces the same result as ``p=2`` in all cases</span>
<span class="sd">            except when :attr:`dim` is a list of three or more dims, in which</span>
<span class="sd">            case Frobenius norm throws an error.</span>

<span class="sd">            Nuclear norm can only be calculated across exactly two dimensions.</span>

<span class="sd">        dim (int, tuple of ints, list of ints, optional):</span>
<span class="sd">            Specifies which dimension or dimensions of :attr:`input` to</span>
<span class="sd">            calculate the norm across. If :attr:`dim` is ``None``, the norm will</span>
<span class="sd">            be calculated across all dimensions of :attr:`input`. If the norm</span>
<span class="sd">            type indicated by :attr:`p` does not support the specified number of</span>
<span class="sd">            dimensions, an error will occur.</span>
<span class="sd">        keepdim (bool, optional): whether the output tensors have :attr:`dim`</span>
<span class="sd">            retained or not. Ignored if :attr:`dim` = ``None`` and</span>
<span class="sd">            :attr:`out` = ``None``. Default: ``False``</span>
<span class="sd">        out (Tensor, optional): the output tensor. Ignored if</span>
<span class="sd">            :attr:`dim` = ``None`` and :attr:`out` = ``None``.</span>
<span class="sd">        dtype (:class:`torch.dtype`, optional): the desired data type of</span>
<span class="sd">            returned tensor. If specified, the input tensor is casted to</span>
<span class="sd">            :attr:`dtype` while performing the operation. Default: None.</span>

<span class="sd">    .. note::</span>
<span class="sd">        Even though ``p=&#39;fro&#39;`` supports any number of dimensions, the true</span>
<span class="sd">        mathematical definition of Frobenius norm only applies to tensors with</span>
<span class="sd">        exactly two dimensions. :func:`torch.linalg.matrix_norm` with ``ord=&#39;fro&#39;``</span>
<span class="sd">        aligns with the mathematical definition, since it can only be applied across</span>
<span class="sd">        exactly two dimensions.</span>

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

<span class="sd">        &gt;&gt;&gt; import torch</span>
<span class="sd">        &gt;&gt;&gt; a = torch.arange(9, dtype= torch.float) - 4</span>
<span class="sd">        &gt;&gt;&gt; b = a.reshape((3, 3))</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(a)</span>
<span class="sd">        tensor(7.7460)</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(b)</span>
<span class="sd">        tensor(7.7460)</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(a, float(&#39;inf&#39;))</span>
<span class="sd">        tensor(4.)</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(b, float(&#39;inf&#39;))</span>
<span class="sd">        tensor(4.)</span>
<span class="sd">        &gt;&gt;&gt; c = torch.tensor([[ 1, 2, 3], [-1, 1, 4]] , dtype=torch.float)</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(c, dim=0)</span>
<span class="sd">        tensor([1.4142, 2.2361, 5.0000])</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(c, dim=1)</span>
<span class="sd">        tensor([3.7417, 4.2426])</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(c, p=1, dim=1)</span>
<span class="sd">        tensor([6., 6.])</span>
<span class="sd">        &gt;&gt;&gt; d = torch.arange(8, dtype=torch.float).reshape(2, 2, 2)</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(d, dim=(1, 2))</span>
<span class="sd">        tensor([ 3.7417, 11.2250])</span>
<span class="sd">        &gt;&gt;&gt; torch.norm(d[0, :, :]), torch.norm(d[1, :, :])</span>
<span class="sd">        (tensor(3.7417), tensor(11.2250))</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">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">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">keepdim</span><span class="o">=</span><span class="n">keepdim</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="n">dtype</span><span class="o">=</span><span class="n">dtype</span><span class="p">)</span>

    <span class="c1"># NB. All the repeated code and weird python is to please TorchScript.</span>
    <span class="c1">#     For a more compact implementation see the relevant function in `_refs/__init__.py`</span>

    <span class="c1"># We don&#39;t do this for MPS or sparse tensors</span>
    <span class="k">if</span> <span class="nb">input</span><span class="o">.</span><span class="n">layout</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">strided</span> <span class="ow">and</span> <span class="nb">input</span><span class="o">.</span><span class="n">device</span><span class="o">.</span><span class="n">type</span> <span class="ow">in</span> \
            <span class="p">(</span><span class="s2">&quot;cpu&quot;</span><span class="p">,</span> <span class="s2">&quot;cuda&quot;</span><span class="p">,</span> <span class="s2">&quot;meta&quot;</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">backend_registration</span><span class="o">.</span><span class="n">_privateuse1_backend_name</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">dim</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="nb">isinstance</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
                <span class="n">_dim</span> <span class="o">=</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">_dim</span> <span class="o">=</span> <span class="n">dim</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">_dim</span> <span class="o">=</span> <span class="kc">None</span>  <span class="c1"># type: ignore[assignment]</span>

        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">p</span> <span class="o">==</span> <span class="s2">&quot;fro&quot;</span> <span class="ow">and</span> <span class="p">(</span><span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">))</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</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">if</span> <span class="n">out</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linalg</span><span class="o">.</span><span class="n">vector_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">_dim</span><span class="p">,</span> <span class="n">keepdim</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">else</span><span class="p">:</span>
                    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linalg</span><span class="o">.</span><span class="n">vector_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">_dim</span><span class="p">,</span> <span class="n">keepdim</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="n">out</span><span class="o">=</span><span class="n">out</span><span class="p">)</span>

            <span class="c1"># Here we either call the nuclear norm, or we call matrix_norm with some arguments</span>
            <span class="c1"># that will throw an error</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="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="nb">input</span><span class="o">.</span><span class="n">ndim</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="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linalg</span><span class="o">.</span><span class="n">matrix_norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</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">else</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linalg</span><span class="o">.</span><span class="n">matrix_norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</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="n">out</span><span class="o">=</span><span class="n">out</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># NB. p should be Union[str, number], not Optional!</span>
            <span class="n">_p</span> <span class="o">=</span> <span class="mf">2.0</span> <span class="k">if</span> <span class="n">p</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">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="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linalg</span><span class="o">.</span><span class="n">vector_norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</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">else</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">linalg</span><span class="o">.</span><span class="n">vector_norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</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="n">out</span><span class="o">=</span><span class="n">out</span><span class="p">)</span>

    <span class="n">ndim</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="c1"># catch default case</span>
    <span class="k">if</span> <span class="n">dim</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">out</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">dtype</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">p</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="nb">isinstance</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">p</span> <span class="o">==</span> <span class="s2">&quot;fro&quot;</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">frobenius_norm</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="p">(),</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
            <span class="n">_dim</span> <span class="o">=</span> <span class="p">[</span><span class="n">i</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">ndim</span><span class="p">)]</span>  <span class="c1"># noqa: C416 TODO: rewrite as list(range(m))</span>
            <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">norm</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">dim</span><span class="o">=</span><span class="n">_dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>

    <span class="c1"># TODO: when https://github.com/pytorch/pytorch/issues/33782 is fixed</span>
    <span class="c1"># remove the overloads where dim is an int and replace with BraodcastingList1</span>
    <span class="c1"># and remove next four lines, replace _dim with dim</span>
    <span class="k">if</span> <span class="n">dim</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="nb">isinstance</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
            <span class="n">_dim</span> <span class="o">=</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">_dim</span> <span class="o">=</span> <span class="n">dim</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">_dim</span> <span class="o">=</span> <span class="kc">None</span>  <span class="c1"># type: ignore[assignment]</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">p</span> <span class="o">==</span> <span class="s2">&quot;fro&quot;</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">dtype</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;dtype argument is not supported in frobenius norm&quot;</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="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">ndim</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="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">frobenius_norm</span><span class="p">(</span><span class="nb">input</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="n">keepdim</span><span class="p">)</span>  <span class="c1"># type: ignore[arg-type]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">frobenius_norm</span><span class="p">(</span><span class="nb">input</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="n">keepdim</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="c1"># type: ignore[arg-type]</span>
        <span class="k">elif</span> <span class="n">p</span> <span class="o">==</span> <span class="s2">&quot;nuc&quot;</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">dtype</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;dtype argument is not supported in nuclear norm&quot;</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="k">if</span> <span class="n">out</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">nuclear_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</span><span class="p">)</span>  <span class="c1"># type: ignore[arg-type]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">nuclear_norm</span><span class="p">(</span><span class="nb">input</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</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="c1"># type: ignore[arg-type]</span>
            <span class="k">else</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="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">nuclear_norm</span><span class="p">(</span><span class="nb">input</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="n">keepdim</span><span class="p">)</span>  <span class="c1"># type: ignore[arg-type]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">nuclear_norm</span><span class="p">(</span><span class="nb">input</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="n">keepdim</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="c1"># type: ignore[arg-type]</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;only valid string values are &#39;fro&#39; and &#39;nuc&#39;, found </span><span class="si">{</span><span class="n">p</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">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="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">ndim</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="k">if</span> <span class="n">dtype</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</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="c1"># type: ignore[attr-defined]</span>
        <span class="k">else</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="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</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="c1"># type: ignore[attr-defined]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">norm</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">_dim</span><span class="p">,</span> <span class="n">keepdim</span><span class="o">=</span><span class="n">keepdim</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="n">out</span><span class="o">=</span><span class="n">out</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span></div>

<div class="viewcode-block" id="unravel_index"><a class="viewcode-back" href="../../generated/torch.unravel_index.html#torch.unravel_index">[docs]</a><span class="k">def</span> <span class="nf">unravel_index</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">shape</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="n">Sequence</span><span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="n">torch</span><span class="o">.</span><span class="n">Size</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="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Converts a tensor of flat indices into a tuple of coordinate tensors that</span>
<span class="sd">    index into an arbitrary tensor of the specified shape.</span>

<span class="sd">    Args:</span>
<span class="sd">        indices (Tensor): An integer tensor containing indices into the</span>
<span class="sd">            flattened version of an arbitrary tensor of shape :attr:`shape`.</span>
<span class="sd">            All elements must be in the range ``[0, prod(shape) - 1]``.</span>

<span class="sd">        shape (int, sequence of ints, or torch.Size): The shape of the arbitrary</span>
<span class="sd">            tensor. All elements must be non-negative.</span>

<span class="sd">    Returns:</span>
<span class="sd">        tuple of Tensors: Each ``i``-th tensor in the ouput corresponds with</span>
<span class="sd">        dimension ``i`` of :attr:`shape`. Each tensor has the same shape as</span>
<span class="sd">        ``indices`` and contains one index into dimension ``i`` for each of the</span>
<span class="sd">        flat indices given by ``indices``.</span>

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

<span class="sd">        &gt;&gt;&gt; import torch</span>
<span class="sd">        &gt;&gt;&gt; torch.unravel_index(torch.tensor(4), (3, 2))</span>
<span class="sd">        (tensor(2),</span>
<span class="sd">         tensor(0))</span>

<span class="sd">        &gt;&gt;&gt; torch.unravel_index(torch.tensor([4, 1]), (3, 2))</span>
<span class="sd">        (tensor([2, 0]),</span>
<span class="sd">         tensor([0, 1]))</span>

<span class="sd">        &gt;&gt;&gt; torch.unravel_index(torch.tensor([0, 1, 2, 3, 4, 5]), (3, 2))</span>
<span class="sd">        (tensor([0, 0, 1, 1, 2, 2]),</span>
<span class="sd">         tensor([0, 1, 0, 1, 0, 1]))</span>

<span class="sd">        &gt;&gt;&gt; torch.unravel_index(torch.tensor([1234, 5678]), (10, 10, 10, 10))</span>
<span class="sd">        (tensor([1, 5]),</span>
<span class="sd">         tensor([2, 6]),</span>
<span class="sd">         tensor([3, 7]),</span>
<span class="sd">         tensor([4, 8]))</span>

<span class="sd">        &gt;&gt;&gt; torch.unravel_index(torch.tensor([[1234], [5678]]), (10, 10, 10, 10))</span>
<span class="sd">        (tensor([[1], [5]]),</span>
<span class="sd">         tensor([[2], [6]]),</span>
<span class="sd">         tensor([[3], [7]]),</span>
<span class="sd">         tensor([[4], [8]]))</span>

<span class="sd">        &gt;&gt;&gt; torch.unravel_index(torch.tensor([[1234], [5678]]), (100, 100))</span>
<span class="sd">        (tensor([[12], [56]]),</span>
<span class="sd">         tensor([[34], [78]]))</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">indices</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">unravel_index</span><span class="p">,</span> <span class="p">(</span><span class="n">indices</span><span class="p">,),</span> <span class="n">indices</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="n">shape</span><span class="p">)</span>
    <span class="n">res_tensor</span> <span class="o">=</span> <span class="n">_unravel_index</span><span class="p">(</span><span class="n">indices</span><span class="p">,</span> <span class="n">shape</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">res_tensor</span><span class="o">.</span><span class="n">unbind</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span></div>

<span class="k">def</span> <span class="nf">_unravel_index</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">shape</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="n">Sequence</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="n">torch</span><span class="o">.</span><span class="n">_check_type</span><span class="p">(</span>
        <span class="ow">not</span> <span class="n">indices</span><span class="o">.</span><span class="n">is_complex</span><span class="p">()</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">indices</span><span class="o">.</span><span class="n">is_floating_point</span><span class="p">()</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">indices</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">bool</span><span class="p">,</span>
        <span class="k">lambda</span><span class="p">:</span> <span class="sa">f</span><span class="s2">&quot;expected &#39;indices&#39; to be integer dtype, but got </span><span class="si">{</span><span class="n">indices</span><span class="o">.</span><span class="n">dtype</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

    <span class="n">torch</span><span class="o">.</span><span class="n">_check_type</span><span class="p">(</span>
        <span class="nb">isinstance</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">)),</span>
        <span class="k">lambda</span><span class="p">:</span> <span class="sa">f</span><span class="s2">&quot;expected &#39;shape&#39; to be int or sequence of ints, but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">shape</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)):</span>
        <span class="n">shape</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">Size</span><span class="p">([</span><span class="n">shape</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">dim</span> <span class="ow">in</span> <span class="n">shape</span><span class="p">:</span>
            <span class="n">torch</span><span class="o">.</span><span class="n">_check_type</span><span class="p">(</span>
                <span class="nb">isinstance</span><span class="p">(</span><span class="n">dim</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">SymInt</span><span class="p">)),</span>
                <span class="k">lambda</span><span class="p">:</span> <span class="sa">f</span><span class="s2">&quot;expected &#39;shape&#39; sequence to only contain ints, but got </span><span class="si">{</span><span class="nb">type</span><span class="p">(</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">shape</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">Size</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span>

    <span class="n">torch</span><span class="o">.</span><span class="n">_check_value</span><span class="p">(</span>
        <span class="nb">all</span><span class="p">(</span><span class="n">dim</span> <span class="o">&gt;=</span> <span class="mi">0</span> <span class="k">for</span> <span class="n">dim</span> <span class="ow">in</span> <span class="n">shape</span><span class="p">),</span>
        <span class="k">lambda</span><span class="p">:</span> <span class="sa">f</span><span class="s2">&quot;&#39;shape&#39; cannot have negative values, but got </span><span class="si">{</span><span class="nb">tuple</span><span class="p">(</span><span class="n">shape</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>

    <span class="n">coefs</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">reversed</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">itertools</span><span class="o">.</span><span class="n">accumulate</span><span class="p">(</span><span class="nb">reversed</span><span class="p">(</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span> <span class="o">+</span> <span class="n">torch</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">func</span><span class="o">=</span><span class="n">operator</span><span class="o">.</span><span class="n">mul</span><span class="p">))))</span>
    <span class="k">return</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="o">.</span><span class="n">floor_divide</span><span class="p">(</span>
        <span class="n">torch</span><span class="o">.</span><span class="n">tensor</span><span class="p">(</span><span class="n">coefs</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">indices</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">int64</span><span class="p">)</span>
    <span class="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">shape</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">indices</span><span class="o">.</span><span class="n">device</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">int64</span><span class="p">)</span>

<div class="viewcode-block" id="chain_matmul"><a class="viewcode-back" href="../../generated/torch.chain_matmul.html#torch.chain_matmul">[docs]</a><span class="k">def</span> <span class="nf">chain_matmul</span><span class="p">(</span><span class="o">*</span><span class="n">matrices</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Returns the matrix product of the :math:`N` 2-D tensors. This product is efficiently computed</span>
<span class="sd">    using the matrix chain order algorithm which selects the order in which incurs the lowest cost in terms</span>
<span class="sd">    of arithmetic operations (`[CLRS]`_). Note that since this is a function to compute the product, :math:`N`</span>
<span class="sd">    needs to be greater than or equal to 2; if equal to 2 then a trivial matrix-matrix product is returned.</span>
<span class="sd">    If :math:`N` is 1, then this is a no-op - the original matrix is returned as is.</span>

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

<span class="sd">        :func:`torch.chain_matmul` is deprecated and will be removed in a future PyTorch release.</span>
<span class="sd">        Use :func:`torch.linalg.multi_dot` instead, which accepts a list of two or more tensors</span>
<span class="sd">        rather than multiple arguments.</span>

<span class="sd">    Args:</span>
<span class="sd">        matrices (Tensors...): a sequence of 2 or more 2-D tensors whose product is to be determined.</span>
<span class="sd">        out (Tensor, optional): the output tensor. Ignored if :attr:`out` = ``None``.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Tensor: if the :math:`i^{th}` tensor was of dimensions :math:`p_{i} \times p_{i + 1}`, then the product</span>
<span class="sd">        would be of dimensions :math:`p_{1} \times p_{N + 1}`.</span>

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

<span class="sd">        &gt;&gt;&gt; # xdoctest: +SKIP</span>
<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; a = torch.randn(3, 4)</span>
<span class="sd">        &gt;&gt;&gt; b = torch.randn(4, 5)</span>
<span class="sd">        &gt;&gt;&gt; c = torch.randn(5, 6)</span>
<span class="sd">        &gt;&gt;&gt; d = torch.randn(6, 7)</span>
<span class="sd">        &gt;&gt;&gt; # will raise a deprecation warning</span>
<span class="sd">        &gt;&gt;&gt; torch.chain_matmul(a, b, c, d)</span>
<span class="sd">        tensor([[ -2.3375,  -3.9790,  -4.1119,  -6.6577,   9.5609, -11.5095,  -3.2614],</span>
<span class="sd">                [ 21.4038,   3.3378,  -8.4982,  -5.2457, -10.2561,  -2.4684,   2.7163],</span>
<span class="sd">                [ -0.9647,  -5.8917,  -2.3213,  -5.2284,  12.8615, -12.2816,  -2.5095]])</span>

<span class="sd">    .. _`[CLRS]`: https://mitpress.mit.edu/books/introduction-algorithms-third-edition</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># This wrapper exists to support variadic args.</span>
    <span class="k">if</span> <span class="n">has_torch_function</span><span class="p">(</span><span class="n">matrices</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span><span class="n">chain_matmul</span><span class="p">,</span> <span class="n">matrices</span><span class="p">,</span> <span class="o">*</span><span class="n">matrices</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="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">chain_matmul</span><span class="p">(</span><span class="n">matrices</span><span class="p">)</span>  <span class="c1"># type: ignore[attr-defined]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">_VF</span><span class="o">.</span><span class="n">chain_matmul</span><span class="p">(</span><span class="n">matrices</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="c1"># type: ignore[attr-defined]</span></div>


<span class="k">def</span> <span class="nf">_lu_impl</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">get_infos</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, Any) -&gt; Tuple[Tensor, Tensor, Tensor]</span>
<span class="w">    </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Computes the LU factorization of a matrix or batches of matrices</span>
<span class="sd">    :attr:`A`. Returns a tuple containing the LU factorization and</span>
<span class="sd">    pivots of :attr:`A`.  Pivoting is done if :attr:`pivot` is set to</span>
<span class="sd">    ``True``.</span>

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

<span class="sd">        :func:`torch.lu` is deprecated in favor of :func:`torch.linalg.lu_factor`</span>
<span class="sd">        and :func:`torch.linalg.lu_factor_ex`. :func:`torch.lu` will be removed in a</span>
<span class="sd">        future PyTorch release.</span>
<span class="sd">        ``LU, pivots, info = torch.lu(A, compute_pivots)`` should be replaced with</span>

<span class="sd">        .. code:: python</span>

<span class="sd">            LU, pivots = torch.linalg.lu_factor(A, compute_pivots)</span>

<span class="sd">        ``LU, pivots, info = torch.lu(A, compute_pivots, get_infos=True)`` should be replaced with</span>

<span class="sd">        .. code:: python</span>

<span class="sd">            LU, pivots, info = torch.linalg.lu_factor_ex(A, compute_pivots)</span>

<span class="sd">    .. note::</span>
<span class="sd">        * The returned permutation matrix for every matrix in the batch is</span>
<span class="sd">          represented by a 1-indexed vector of size ``min(A.shape[-2], A.shape[-1])``.</span>
<span class="sd">          ``pivots[i] == j`` represents that in the ``i``-th step of the algorithm,</span>
<span class="sd">          the ``i``-th row was permuted with the ``j-1``-th row.</span>
<span class="sd">        * LU factorization with :attr:`pivot` = ``False`` is not available</span>
<span class="sd">          for CPU, and attempting to do so will throw an error. However,</span>
<span class="sd">          LU factorization with :attr:`pivot` = ``False`` is available for</span>
<span class="sd">          CUDA.</span>
<span class="sd">        * This function does not check if the factorization was successful</span>
<span class="sd">          or not if :attr:`get_infos` is ``True`` since the status of the</span>
<span class="sd">          factorization is present in the third element of the return tuple.</span>
<span class="sd">        * In the case of batches of square matrices with size less or equal</span>
<span class="sd">          to 32 on a CUDA device, the LU factorization is repeated for</span>
<span class="sd">          singular matrices due to the bug in the MAGMA library</span>
<span class="sd">          (see magma issue 13).</span>
<span class="sd">        * ``L``, ``U``, and ``P`` can be derived using :func:`torch.lu_unpack`.</span>

<span class="sd">    .. warning::</span>
<span class="sd">        The gradients of this function will only be finite when :attr:`A` is full rank.</span>
<span class="sd">        This is because the LU decomposition is just differentiable at full rank matrices.</span>
<span class="sd">        Furthermore, if :attr:`A` is close to not being full rank,</span>
<span class="sd">        the gradient will be numerically unstable as it depends on the computation of :math:`L^{-1}` and :math:`U^{-1}`.</span>

<span class="sd">    Args:</span>
<span class="sd">        A (Tensor): the tensor to factor of size :math:`(*, m, n)`</span>
<span class="sd">        pivot (bool, optional): controls whether pivoting is done. Default: ``True``</span>
<span class="sd">        get_infos (bool, optional): if set to ``True``, returns an info IntTensor.</span>
<span class="sd">                                    Default: ``False``</span>
<span class="sd">        out (tuple, optional): optional output tuple. If :attr:`get_infos` is ``True``,</span>
<span class="sd">                               then the elements in the tuple are Tensor, IntTensor,</span>
<span class="sd">                               and IntTensor. If :attr:`get_infos` is ``False``, then the</span>
<span class="sd">                               elements in the tuple are Tensor, IntTensor. Default: ``None``</span>

<span class="sd">    Returns:</span>
<span class="sd">        (Tensor, IntTensor, IntTensor (optional)): A tuple of tensors containing</span>

<span class="sd">            - **factorization** (*Tensor*): the factorization of size :math:`(*, m, n)`</span>

<span class="sd">            - **pivots** (*IntTensor*): the pivots of size :math:`(*, \text{min}(m, n))`.</span>
<span class="sd">              ``pivots`` stores all the intermediate transpositions of rows.</span>
<span class="sd">              The final permutation ``perm`` could be reconstructed by</span>
<span class="sd">              applying ``swap(perm[i], perm[pivots[i] - 1])`` for ``i = 0, ..., pivots.size(-1) - 1``,</span>
<span class="sd">              where ``perm`` is initially the identity permutation of :math:`m` elements</span>
<span class="sd">              (essentially this is what :func:`torch.lu_unpack` is doing).</span>

<span class="sd">            - **infos** (*IntTensor*, *optional*): if :attr:`get_infos` is ``True``, this is a tensor of</span>
<span class="sd">              size :math:`(*)` where non-zero values indicate whether factorization for the matrix or</span>
<span class="sd">              each minibatch has succeeded or failed</span>

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

<span class="sd">        &gt;&gt;&gt; # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)</span>
<span class="sd">        &gt;&gt;&gt; # xdoctest: +IGNORE_WANT(&quot;non-deterministic&quot;)</span>
<span class="sd">        &gt;&gt;&gt; A = torch.randn(2, 3, 3)</span>
<span class="sd">        &gt;&gt;&gt; A_LU, pivots = torch.lu(A)</span>
<span class="sd">        &gt;&gt;&gt; A_LU</span>
<span class="sd">        tensor([[[ 1.3506,  2.5558, -0.0816],</span>
<span class="sd">                 [ 0.1684,  1.1551,  0.1940],</span>
<span class="sd">                 [ 0.1193,  0.6189, -0.5497]],</span>

<span class="sd">                [[ 0.4526,  1.2526, -0.3285],</span>
<span class="sd">                 [-0.7988,  0.7175, -0.9701],</span>
<span class="sd">                 [ 0.2634, -0.9255, -0.3459]]])</span>
<span class="sd">        &gt;&gt;&gt; pivots</span>
<span class="sd">        tensor([[ 3,  3,  3],</span>
<span class="sd">                [ 3,  3,  3]], dtype=torch.int32)</span>
<span class="sd">        &gt;&gt;&gt; A_LU, pivots, info = torch.lu(A, get_infos=True)</span>
<span class="sd">        &gt;&gt;&gt; if info.nonzero().size(0) == 0:</span>
<span class="sd">        ...     print(&#39;LU factorization succeeded for all samples!&#39;)</span>
<span class="sd">        LU factorization succeeded for all samples!</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># If get_infos is True, then we don&#39;t need to check for errors and vice versa</span>
    <span class="k">return</span> <span class="n">torch</span><span class="o">.</span><span class="n">_lu_with_info</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="o">=</span><span class="n">pivot</span><span class="p">,</span> <span class="n">check_errors</span><span class="o">=</span><span class="p">(</span><span class="ow">not</span> <span class="n">get_infos</span><span class="p">))</span>

<span class="k">if</span> <span class="n">TYPE_CHECKING</span><span class="p">:</span>
    <span class="n">_ListOrSeq</span> <span class="o">=</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span>
<span class="k">else</span><span class="p">:</span>
    <span class="n">_ListOrSeq</span> <span class="o">=</span> <span class="n">List</span><span class="p">[</span><span class="n">Tensor</span><span class="p">]</span>


<span class="k">def</span> <span class="nf">_check_list_size</span><span class="p">(</span><span class="n">out_len</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">get_infos</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> <span class="n">out</span><span class="p">:</span> <span class="n">_ListOrSeq</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
    <span class="n">get_infos_int</span> <span class="o">=</span> <span class="mi">1</span> <span class="k">if</span> <span class="n">get_infos</span> <span class="k">else</span> <span class="mi">0</span>
    <span class="k">if</span> <span class="n">out_len</span> <span class="o">-</span> <span class="n">get_infos_int</span> <span class="o">!=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;expected tuple of </span><span class="si">{</span><span class="mi">2</span><span class="w"> </span><span class="o">+</span><span class="w"> </span><span class="nb">int</span><span class="p">(</span><span class="n">get_infos</span><span class="p">)</span><span class="si">}</span><span class="s2"> elements but got </span><span class="si">{</span><span class="n">out_len</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">out</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;argument &#39;out&#39; must be tuple of Tensors, not </span><span class="si">{</span><span class="nb">type</span><span class="p">(</span><span class="n">out</span><span class="p">)</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_lu_with_infos</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">get_infos</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, Optional[Tuple[Tensor, Tensor, Tensor]]) -&gt; Tuple[Tensor, Tensor, Tensor]</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="n">A</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">lu</span><span class="p">,</span> <span class="p">(</span><span class="n">A</span><span class="p">,),</span> <span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="o">=</span><span class="n">pivot</span><span class="p">,</span> <span class="n">get_infos</span><span class="o">=</span><span class="n">get_infos</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="n">result</span> <span class="o">=</span> <span class="n">_lu_impl</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="p">,</span> <span class="n">get_infos</span><span class="p">,</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="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">_check_list_size</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">out</span><span class="p">),</span> <span class="n">get_infos</span><span class="p">,</span> <span class="n">out</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">out</span><span class="p">)):</span>
            <span class="n">out</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">resize_as_</span><span class="p">(</span><span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">])</span><span class="o">.</span><span class="n">copy_</span><span class="p">(</span><span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">out</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">result</span>  <span class="c1"># A_LU, pivots, infos</span>


<span class="k">def</span> <span class="nf">_lu_no_infos</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">get_infos</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="c1"># type: (Tensor, bool, bool, Optional[Tuple[Tensor, Tensor]]) -&gt; Tuple[Tensor, Tensor]</span>
    <span class="c1"># need to check for torch_function here so that we exit if</span>
    <span class="k">if</span> <span class="n">has_torch_function_unary</span><span class="p">(</span><span class="n">A</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">handle_torch_function</span><span class="p">(</span>
            <span class="n">lu</span><span class="p">,</span> <span class="p">(</span><span class="n">A</span><span class="p">,),</span> <span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="o">=</span><span class="n">pivot</span><span class="p">,</span> <span class="n">get_infos</span><span class="o">=</span><span class="n">get_infos</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="n">result</span> <span class="o">=</span> <span class="n">_lu_impl</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">pivot</span><span class="p">,</span> <span class="n">get_infos</span><span class="p">,</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="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">_check_list_size</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">out</span><span class="p">),</span> <span class="n">get_infos</span><span class="p">,</span> <span class="n">out</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">out</span><span class="p">)):</span>
            <span class="n">out</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="o">.</span><span class="n">resize_as_</span><span class="p">(</span><span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">])</span><span class="o">.</span><span class="n">copy_</span><span class="p">(</span><span class="n">result</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="k">return</span> <span class="n">out</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">result</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">result</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>  <span class="c1"># A_LU, pivots</span>

<span class="c1"># The return type of lu depends on `get_infos`, so in order to resolve the output type</span>
<span class="c1"># of lu in TorchScript we need to statically know the value of `get_infos`</span>
<span class="n">lu</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="s1">&#39;get_infos&#39;</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">_lu_with_infos</span><span class="p">,</span>
    <span class="n">if_false</span><span class="o">=</span><span class="n">_lu_no_infos</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="s1">&#39;lu&#39;</span><span class="p">)</span>
<span class="n">lu</span><span class="o">.</span><span class="vm">__doc__</span> <span class="o">=</span> <span class="n">_lu_impl</span><span class="o">.</span><span class="vm">__doc__</span>


<span class="k">def</span> <span class="nf">align_tensors</span><span class="p">(</span><span class="o">*</span><span class="n">tensors</span><span class="p">):</span>
    <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s1">&#39;`align_tensors` not yet implemented.&#39;</span><span class="p">)</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/sphinx_highlight.js"></script>
         <script src="../../_static/clipboard.min.js"></script>
         <script src="../../_static/copybutton.js"></script>
     

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

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

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


  <!-- Begin Footer -->

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

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

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

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

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

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

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

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

  </footer>

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

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

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

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>
            
          <li>
            <a href="">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>