|
235 | 235 | <div class="pytorch-left-menu-search"> |
236 | 236 |
|
237 | 237 | <div class="version"> |
238 | | - <a href='https://pytorch.org/docs/versions.html'>master (2.0.0a0+gitceb4435 ) ▼</a> |
| 238 | + <a href='https://pytorch.org/docs/versions.html'>master (2.0.0a0+git1ff292a ) ▼</a> |
239 | 239 | </div> |
240 | 240 |
|
241 | 241 |
|
@@ -1058,8 +1058,7 @@ <h1>Source code for torch._dynamo.eval_frame</h1><div class="highlight"><pre> |
1058 | 1058 |
|
1059 | 1059 | <span class="k">return</span> <span class="n">result_capturing_wrapper</span> |
1060 | 1060 |
|
1061 | | - <span class="c1"># TODO(voz): Handle kwargs properly?</span> |
1062 | | - <span class="n">flat_args</span><span class="p">,</span> <span class="n">in_spec</span> <span class="o">=</span> <span class="n">pytree</span><span class="o">.</span><span class="n">tree_flatten</span><span class="p">(</span><span class="n">args</span><span class="p">)</span> |
| 1061 | + <span class="n">flat_args</span><span class="p">,</span> <span class="n">in_spec</span> <span class="o">=</span> <span class="n">pytree</span><span class="o">.</span><span class="n">tree_flatten</span><span class="p">((</span><span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">))</span> |
1063 | 1062 |
|
1064 | 1063 | <span class="n">remove_from_cache</span><span class="p">(</span><span class="n">f</span><span class="p">)</span> |
1065 | 1064 | <span class="k">with</span> <span class="n">patch</span><span class="p">(</span><span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">.most_recent_backend"</span><span class="p">,</span> <span class="kc">None</span><span class="p">):</span> |
@@ -1133,9 +1132,10 @@ <h1>Source code for torch._dynamo.eval_frame</h1><div class="highlight"><pre> |
1133 | 1132 | <span class="p">)</span><span class="o">.</span><span class="n">transform</span><span class="p">()</span> |
1134 | 1133 |
|
1135 | 1134 | <span class="c1"># Make dynamo graph to have same input/output spec as user code</span> |
| 1135 | + <span class="n">input_strs</span> <span class="o">=</span> <span class="p">[</span><span class="sa">f</span><span class="s2">"orig_arg_</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">"</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">args</span><span class="p">))]</span> <span class="o">+</span> <span class="nb">list</span><span class="p">(</span><span class="n">kwargs</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span> |
1136 | 1136 | <span class="n">new_graph</span><span class="o">.</span><span class="n">graph</span><span class="o">.</span><span class="n">_codegen</span> <span class="o">=</span> <span class="n">_PyTreeCodeGen</span><span class="p">(</span> |
1137 | 1137 | <span class="n">_PyTreeInfo</span><span class="p">(</span> |
1138 | | - <span class="p">[</span><span class="sa">f</span><span class="s2">"orig_arg_</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">"</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">args</span><span class="p">))],</span> |
| 1138 | + <span class="n">input_strs</span><span class="p">,</span> |
1139 | 1139 | <span class="n">in_spec</span><span class="p">,</span> |
1140 | 1140 | <span class="n">out_spec_traced</span><span class="p">,</span> |
1141 | 1141 | <span class="p">)</span> |
|
0 commit comments