Generate Python docs from pytorch/pytorch@0052943

pytorch · Dec 7, 2024 · 236fc3e · 236fc3e
1 parent 1980bef
commit 236fc3e
Show file tree

Hide file tree

Showing 41 changed files with 101 additions and 47 deletions.
diff --git a/nightly/_modules/index.html b/nightly/_modules/index.html
@@ -265,7 +265,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/_modules/torch/_functorch/aot_autograd.html b/nightly/_modules/torch/_functorch/aot_autograd.html
@@ -265,7 +265,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 
@@ -389,7 +389,19 @@ <h1>Source code for torch._functorch.aot_autograd</h1><div class="highlight"><pr
 <span class="kn">import</span> <span class="nn">itertools</span>
 <span class="kn">from</span> <span class="nn">contextlib</span> <span class="kn">import</span> <span class="n">contextmanager</span><span class="p">,</span> <span class="n">nullcontext</span>
 <span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">partial</span><span class="p">,</span> <span class="n">wraps</span>
-<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">Dict</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">NewType</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Tuple</span>
+<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="p">(</span>
+    <span class="n">Any</span><span class="p">,</span>
+    <span class="n">Callable</span><span class="p">,</span>
+    <span class="n">Dict</span><span class="p">,</span>
+    <span class="n">List</span><span class="p">,</span>
+    <span class="n">NewType</span><span class="p">,</span>
+    <span class="n">Optional</span><span class="p">,</span>
+    <span class="n">Protocol</span><span class="p">,</span>
+    <span class="n">Sequence</span><span class="p">,</span>
+    <span class="n">Tuple</span><span class="p">,</span>
+    <span class="n">Type</span><span class="p">,</span>
+    <span class="n">TypeVar</span><span class="p">,</span>
+<span class="p">)</span>
 <span class="kn">from</span> <span class="nn">unittest.mock</span> <span class="kn">import</span> <span class="n">patch</span>
 
 <span class="kn">import</span> <span class="nn">torch</span>
@@ -407,7 +419,8 @@ <h1>Source code for torch._functorch.aot_autograd</h1><div class="highlight"><pr
     <span class="n">preserve_rng_state</span><span class="p">,</span>
 <span class="p">)</span>
 <span class="kn">from</span> <span class="nn">torch._guards</span> <span class="kn">import</span> <span class="n">detect_fake_mode</span>
-<span class="kn">from</span> <span class="nn">torch._inductor.utils</span> <span class="kn">import</span> <span class="n">BoxedBool</span>
+<span class="kn">from</span> <span class="nn">torch._inductor.output_code</span> <span class="kn">import</span> <span class="n">OutputCode</span>
+<span class="kn">from</span> <span class="nn">torch._inductor.utils</span> <span class="kn">import</span> <span class="n">BoxedBool</span><span class="p">,</span> <span class="n">InputType</span>
 <span class="kn">from</span> <span class="nn">torch._subclasses</span> <span class="kn">import</span> <span class="n">FakeTensor</span><span class="p">,</span> <span class="n">FakeTensorMode</span>
 <span class="kn">from</span> <span class="nn">torch.fx.experimental.proxy_tensor</span> <span class="kn">import</span> <span class="n">make_fx</span>
 <span class="kn">from</span> <span class="nn">torch.fx.experimental.symbolic_shapes</span> <span class="kn">import</span> <span class="n">ShapeEnv</span>
@@ -821,6 +834,47 @@ <h1>Source code for torch._functorch.aot_autograd</h1><div class="highlight"><pr
 <span class="n">FakifiedFlatArgs</span> <span class="o">=</span> <span class="n">NewType</span><span class="p">(</span><span class="s2">&quot;FakifiedFlatArgs&quot;</span><span class="p">,</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">])</span>
 
 
+<span class="n">TOutputCode</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s2">&quot;TOutputCode&quot;</span><span class="p">,</span> <span class="n">bound</span><span class="o">=</span><span class="n">OutputCode</span><span class="p">)</span>
+
+
+<span class="k">class</span> <span class="nc">AOTDispatchCompiler</span><span class="p">(</span><span class="n">Protocol</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Represents a fw or bw_compiler passed to AOTAutograd.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span>
+        <span class="n">example_inputs</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">InputType</span><span class="p">],</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
+        <span class="o">...</span>
+
+
+<span class="c1"># TODO: bikeshed on this name</span>
+<span class="k">class</span> <span class="nc">SerializableAOTDispatchCompiler</span><span class="p">(</span><span class="n">AOTDispatchCompiler</span><span class="p">):</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Represents an AOTDispatchCompiler that returns an OutputCode, and is</span>
+<span class="sd">    therefore cacheable. SerializableAOTDispatchCompiler always return an OutputCode.</span>
+<span class="sd">    A _CompileFxCallable usually gets converted into an AOTDispatchCompiler after binding all of</span>
+<span class="sd">    the kwargs in _CompileFxKwargs.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">output_code_ty</span><span class="p">:</span> <span class="n">Type</span><span class="p">[</span><span class="n">TOutputCode</span><span class="p">],</span>
+        <span class="n">compiler_fn</span><span class="p">:</span> <span class="n">Callable</span><span class="p">[[</span><span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">InputType</span><span class="p">]],</span> <span class="n">TOutputCode</span><span class="p">],</span>
+    <span class="p">):</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">output_code_ty</span> <span class="o">=</span> <span class="n">output_code_ty</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">compiler_fn</span> <span class="o">=</span> <span class="n">compiler_fn</span>
+
+    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span>
+        <span class="bp">self</span><span class="p">,</span>
+        <span class="n">gm</span><span class="p">:</span> <span class="n">torch</span><span class="o">.</span><span class="n">fx</span><span class="o">.</span><span class="n">GraphModule</span><span class="p">,</span>
+        <span class="n">example_inputs</span><span class="p">:</span> <span class="n">Sequence</span><span class="p">[</span><span class="n">InputType</span><span class="p">],</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">OutputCode</span><span class="p">:</span>
+        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">compiler_fn</span><span class="p">(</span><span class="n">gm</span><span class="p">,</span> <span class="n">example_inputs</span><span class="p">)</span>
+
+
 <span class="k">def</span> <span class="nf">process_inputs</span><span class="p">(</span>
     <span class="n">flat_args</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Any</span><span class="p">],</span>
     <span class="n">aot_config</span><span class="p">:</span> <span class="n">AOTConfig</span><span class="p">,</span>
@@ -1339,12 +1393,12 @@ <h1>Source code for torch._functorch.aot_autograd</h1><div class="highlight"><pr
 <span class="k">def</span> <span class="nf">aot_module_simplified</span><span class="p">(</span>
     <span class="n">mod</span><span class="p">:</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">,</span>
     <span class="n">args</span><span class="p">,</span>
-    <span class="n">fw_compiler</span><span class="p">:</span> <span class="n">Callable</span><span class="p">,</span>
-    <span class="n">bw_compiler</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Callable</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
+    <span class="n">fw_compiler</span><span class="p">:</span> <span class="n">AOTDispatchCompiler</span><span class="p">,</span>
+    <span class="n">bw_compiler</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AOTDispatchCompiler</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
     <span class="n">partition_fn</span><span class="p">:</span> <span class="n">Callable</span> <span class="o">=</span> <span class="n">default_partition</span><span class="p">,</span>
     <span class="n">decompositions</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Dict</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
     <span class="n">keep_inference_input_mutations</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
-    <span class="n">inference_compiler</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Callable</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
+    <span class="n">inference_compiler</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">AOTDispatchCompiler</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
     <span class="n">cudagraphs</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">BoxedBool</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
 <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">:</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;</span>
@@ -1472,8 +1526,8 @@ <h1>Source code for torch._functorch.aot_autograd</h1><div class="highlight"><pr
     <span class="c1"># Autograd cache stuff</span>
     <span class="n">remote</span> <span class="o">=</span> <span class="n">should_use_remote_autograd_cache</span><span class="p">()</span>
     <span class="n">local</span> <span class="o">=</span> <span class="n">should_use_local_autograd_cache</span><span class="p">()</span>
-
-    <span class="k">if</span> <span class="n">local</span> <span class="ow">or</span> <span class="n">remote</span><span class="p">:</span>
+    <span class="c1"># We only care if the forward will return an OutputCode.</span>
+    <span class="k">if</span> <span class="p">(</span><span class="n">local</span> <span class="ow">or</span> <span class="n">remote</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fw_compiler</span><span class="p">,</span> <span class="n">SerializableAOTDispatchCompiler</span><span class="p">):</span>
         <span class="n">compiled_fn</span> <span class="o">=</span> <span class="n">AOTAutogradCache</span><span class="o">.</span><span class="n">load</span><span class="p">(</span>
             <span class="n">dispatch_and_compile</span><span class="p">,</span>
             <span class="n">mod</span><span class="p">,</span>

diff --git a/nightly/_modules/torch/_functorch/compilers.html b/nightly/_modules/torch/_functorch/compilers.html
@@ -265,7 +265,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/_modules/torch/_functorch/deprecated.html b/nightly/_modules/torch/_functorch/deprecated.html
@@ -265,7 +265,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/_modules/torch/_functorch/partitioners.html b/nightly/_modules/torch/_functorch/partitioners.html
@@ -265,7 +265,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/aot_autograd.html b/nightly/aot_autograd.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/batch_norm.html b/nightly/batch_norm.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/experimental.html b/nightly/experimental.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/functorch.html b/nightly/functorch.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.combine_state_for_ensemble.html b/nightly/generated/functorch.combine_state_for_ensemble.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.aot_function.html b/nightly/generated/functorch.compile.aot_function.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.aot_module.html b/nightly/generated/functorch.compile.aot_module.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.default_partition.html b/nightly/generated/functorch.compile.default_partition.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.memory_efficient_fusion.html b/nightly/generated/functorch.compile.memory_efficient_fusion.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.min_cut_rematerialization_partition.html b/nightly/generated/functorch.compile.min_cut_rematerialization_partition.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.nop.html b/nightly/generated/functorch.compile.nop.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.compile.ts_compile.html b/nightly/generated/functorch.compile.ts_compile.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.functionalize.html b/nightly/generated/functorch.functionalize.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.grad.html b/nightly/generated/functorch.grad.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.grad_and_value.html b/nightly/generated/functorch.grad_and_value.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.hessian.html b/nightly/generated/functorch.hessian.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.jacfwd.html b/nightly/generated/functorch.jacfwd.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.jacrev.html b/nightly/generated/functorch.jacrev.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>
 
 

diff --git a/nightly/generated/functorch.jvp.html b/nightly/generated/functorch.jvp.html
@@ -267,7 +267,7 @@
           <div class="pytorch-left-menu-search">
 
     <div class="version">
-      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+gitefab8c4) &#x25BC</a>
+      <a href='https://pytorch.org/functorch/versions.html'>nightly (2.6.0a0+git0052943) &#x25BC</a>
     </div>