Update doc from commit 2e1ebfc

master/_modules/index.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/core/xla_model.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/debug/metrics.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/distributed/parallel_loader.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/distributed/spmd/xla_sharding.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/distributed/xla_multiprocessing.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/experimental/eager.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/runtime.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/_modules/torch_xla/torch_xla.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 
@@ -383,8 +383,10 @@
              <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
 
   <h1>Source code for torch_xla.torch_xla</h1><div class="highlight"><pre>
-<span></span><span class="kn">import</span> <span class="nn">contextlib</span>
-<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Tuple</span>
+<span></span><span class="kn">import</span> <span class="nn">collections</span>
+<span class="kn">import</span> <span class="nn">contextlib</span>
+<span class="kn">import</span> <span class="nn">functools</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">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Tuple</span>
 
 <span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">import</span> <span class="nn">torch.distributed</span> <span class="k">as</span> <span class="nn">dist</span>
@@ -444,24 +446,28 @@ <h1>Source code for torch_xla.torch_xla</h1><div class="highlight"><pre>
   <span class="n">xm</span><span class="o">.</span><span class="n">mark_step</span><span class="p">()</span></div>
 
 
-<div class="viewcode-block" id="step"><a class="viewcode-back" href="../../index.html#torch_xla.step">[docs]</a><span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
-<span class="k">def</span> <span class="nf">step</span><span class="p">():</span>
+<div class="viewcode-block" id="step"><a class="viewcode-back" href="../../index.html#torch_xla.step">[docs]</a><span class="k">def</span> <span class="nf">step</span><span class="p">(</span><span class="n">f</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="w">  </span><span class="sd">&quot;&quot;&quot;Wraps code that should be dispatched to the runtime.</span>
 
 <span class="sd">  Experimental: `xla.step` is still a work in progress. Some code that currently</span>
 <span class="sd">  works with `xla.step` but does not follow best practices will become errors in</span>
 <span class="sd">  future releases. See https://github.com/pytorch/xla/issues/6751 for context.</span>
 <span class="sd">  &quot;&quot;&quot;</span>
-  <span class="n">saved_eager_mode_status</span> <span class="o">=</span> <span class="n">torch_xla</span><span class="o">.</span><span class="n">_XLAC</span><span class="o">.</span><span class="n">_get_use_eager_mode</span><span class="p">()</span>
-  <span class="n">torch_xla</span><span class="o">.</span><span class="n">_XLAC</span><span class="o">.</span><span class="n">_set_use_eager_mode</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
-  <span class="c1"># Clear pending operations</span>
-  <span class="n">sync</span><span class="p">()</span>
-
-  <span class="k">try</span><span class="p">:</span>
-    <span class="k">yield</span>
-  <span class="k">finally</span><span class="p">:</span>
+
+  <span class="nd">@contextlib</span><span class="o">.</span><span class="n">contextmanager</span>
+  <span class="k">def</span> <span class="nf">_step</span><span class="p">():</span>
+    <span class="n">saved_eager_mode_status</span> <span class="o">=</span> <span class="n">torch_xla</span><span class="o">.</span><span class="n">_XLAC</span><span class="o">.</span><span class="n">_get_use_eager_mode</span><span class="p">()</span>
+    <span class="n">torch_xla</span><span class="o">.</span><span class="n">_XLAC</span><span class="o">.</span><span class="n">_set_use_eager_mode</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
+    <span class="c1"># Clear pending operations</span>
     <span class="n">sync</span><span class="p">()</span>
-    <span class="n">torch_xla</span><span class="o">.</span><span class="n">_XLAC</span><span class="o">.</span><span class="n">_set_use_eager_mode</span><span class="p">(</span><span class="n">saved_eager_mode_status</span><span class="p">)</span></div>
+
+    <span class="k">try</span><span class="p">:</span>
+      <span class="k">yield</span>
+    <span class="k">finally</span><span class="p">:</span>
+      <span class="n">sync</span><span class="p">()</span>
+      <span class="n">torch_xla</span><span class="o">.</span><span class="n">_XLAC</span><span class="o">.</span><span class="n">_set_use_eager_mode</span><span class="p">(</span><span class="n">saved_eager_mode_status</span><span class="p">)</span>
+
+  <span class="k">return</span> <span class="n">_step</span><span class="p">()</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">f</span> <span class="k">else</span> <span class="n">_step</span><span class="p">()(</span><span class="n">f</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="manual_seed"><a class="viewcode-back" href="../../index.html#torch_xla.manual_seed">[docs]</a><span class="k">def</span> <span class="nf">manual_seed</span><span class="p">(</span><span class="n">seed</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>

master/debug.html

@@ -267,7 +267,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/eager_mode.html

@@ -267,7 +267,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/genindex.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/gpu.html

@@ -267,7 +267,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/index.html

@@ -266,7 +266,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 
@@ -756,7 +756,7 @@ <h1>PyTorch/XLA API<a class="headerlink" href="#pytorch-xla-api" title="Permalin
 
 <dl class="py function">
 <dt class="sig sig-object py" id="torch_xla.step">
-<span class="sig-prename descclassname"><span class="pre">torch_xla.</span></span><span class="sig-name descname"><span class="pre">step</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch_xla/torch_xla.html#step"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch_xla.step" title="Permalink to this definition">¶</a></dt>
+<span class="sig-prename descclassname"><span class="pre">torch_xla.</span></span><span class="sig-name descname"><span class="pre">step</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">f</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Callable</span><span class="p"><span class="pre">]</span></span></span><span class="w"> </span><span class="o"><span class="pre">=</span></span><span class="w"> </span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch_xla/torch_xla.html#step"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch_xla.step" title="Permalink to this definition">¶</a></dt>
 <dd><p>Wraps code that should be dispatched to the runtime.</p>
 <p>Experimental: <cite>xla.step</cite> is still a work in progress. Some code that currently
 works with <cite>xla.step</cite> but does not follow best practices will become errors in

master/multi_process_distributed.html

@@ -267,7 +267,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/notes/source_of_recompilation.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/py-modindex.html

@@ -268,7 +268,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/runtime.html

@@ -267,7 +267,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/search.html

@@ -265,7 +265,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/spmd.html

@@ -267,7 +267,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>
 
 

master/torch_compile.html

@@ -266,7 +266,7 @@
 
 
                 <div class="version">
-                  master (2.5.0+gitd77bf6f )
+                  master (2.5.0+git2e1ebfc )
                 </div>