|
219 | 219 | <div class="pytorch-left-menu-search">
|
220 | 220 |
|
221 | 221 | <div class="version">
|
222 |
| - <a href='https://pytorch.org/docs/versions.html'>master (1.13.0a0+git1910c58 ) ▼</a> |
| 222 | + <a href='https://pytorch.org/docs/versions.html'>master (1.13.0a0+git539076e ) ▼</a> |
223 | 223 | </div>
|
224 | 224 |
|
225 | 225 |
|
|
430 | 430 |
|
431 | 431 | <h1>Source code for torch._tensor_str</h1><div class="highlight"><pre>
|
432 | 432 | <span></span><span class="kn">import</span> <span class="nn">math</span>
|
| 433 | +<span class="kn">import</span> <span class="nn">textwrap</span> |
433 | 434 | <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span>
|
434 | 435 |
|
435 | 436 | <span class="kn">import</span> <span class="nn">torch</span>
|
@@ -795,6 +796,8 @@ <h1>Source code for torch._tensor_str</h1><div class="highlight"><pre>
|
795 | 796 |
|
796 | 797 |
|
797 | 798 | <span class="k">def</span> <span class="nf">_str_intern</span><span class="p">(</span><span class="n">inp</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">tensor_contents</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
|
| 799 | + <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">is_functorch_wrapped_tensor</span><span class="p">(</span><span class="n">inp</span><span class="p">):</span> |
| 800 | + <span class="k">return</span> <span class="n">_functorch_wrapper_str_intern</span><span class="p">(</span><span class="n">inp</span><span class="p">,</span> <span class="n">tensor_contents</span><span class="o">=</span><span class="n">tensor_contents</span><span class="p">)</span> |
798 | 801 | <span class="n">is_plain_tensor</span> <span class="o">=</span> <span class="nb">type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="ow">is</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span> <span class="ow">or</span> <span class="nb">type</span><span class="p">(</span><span class="n">inp</span><span class="p">)</span> <span class="ow">is</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">Parameter</span>
|
799 | 802 | <span class="k">if</span> <span class="n">inp</span><span class="o">.</span><span class="n">is_nested</span><span class="p">:</span>
|
800 | 803 | <span class="n">prefix</span> <span class="o">=</span> <span class="s2">"nested_tensor("</span>
|
@@ -1028,6 +1031,48 @@ <h1>Source code for torch._tensor_str</h1><div class="highlight"><pre>
|
1028 | 1031 | <span class="k">return</span> <span class="n">string_repr</span>
|
1029 | 1032 |
|
1030 | 1033 |
|
| 1034 | +<span class="k">def</span> <span class="nf">_functorch_wrapper_str_intern</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">tensor_contents</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span> |
| 1035 | + <span class="n">level</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">maybe_get_level</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> |
| 1036 | + <span class="k">assert</span> <span class="n">level</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span> |
| 1037 | + |
| 1038 | + <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">is_functionaltensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">):</span> |
| 1039 | + <span class="c1"># Since we're unwrapping the FunctionalTensorWrapper, we need to make sure</span> |
| 1040 | + <span class="c1"># that it's up to date first</span> |
| 1041 | + <span class="n">torch</span><span class="o">.</span><span class="n">_sync</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> |
| 1042 | + |
| 1043 | + <span class="n">value</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">get_unwrapped</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> |
| 1044 | + <span class="n">dl_enabled</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_are_functorch_transforms_active</span><span class="p">()</span> |
| 1045 | + <span class="k">try</span><span class="p">:</span> |
| 1046 | + <span class="c1"># Disable temporarily FuncTorchDynamicLayerFrontMode and</span> |
| 1047 | + <span class="c1"># FuncTorchDynamicLayerBackMode as included dispatch keys</span> |
| 1048 | + <span class="k">if</span> <span class="n">dl_enabled</span><span class="p">:</span> |
| 1049 | + <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">_set_dynamic_layer_keys_included</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span> |
| 1050 | + <span class="n">value_repr</span> <span class="o">=</span> <span class="nb">repr</span><span class="p">(</span><span class="n">value</span><span class="p">)</span> |
| 1051 | + <span class="k">finally</span><span class="p">:</span> |
| 1052 | + <span class="c1"># Reenable FuncTorchDynamicLayerFrontMode and</span> |
| 1053 | + <span class="c1"># FuncTorchDynamicLayerBackMode as included dispatch keys</span> |
| 1054 | + <span class="k">if</span> <span class="n">dl_enabled</span><span class="p">:</span> |
| 1055 | + <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">_set_dynamic_layer_keys_included</span><span class="p">(</span><span class="kc">True</span><span class="p">)</span> |
| 1056 | + |
| 1057 | + <span class="n">indented_value_repr</span> <span class="o">=</span> <span class="n">textwrap</span><span class="o">.</span><span class="n">indent</span><span class="p">(</span><span class="n">value_repr</span><span class="p">,</span> <span class="s2">" "</span> <span class="o">*</span> <span class="mi">4</span><span class="p">)</span> |
| 1058 | + <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">is_batchedtensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">):</span> |
| 1059 | + <span class="n">bdim</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">maybe_get_bdim</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> |
| 1060 | + <span class="k">assert</span> <span class="n">bdim</span> <span class="o">!=</span> <span class="o">-</span><span class="mi">1</span> |
| 1061 | + <span class="k">return</span> <span class="p">(</span> |
| 1062 | + <span class="sa">f</span><span class="s2">"BatchedTensor(lvl=</span><span class="si">{</span><span class="n">level</span><span class="si">}</span><span class="s2">, bdim=</span><span class="si">{</span><span class="n">bdim</span><span class="si">}</span><span class="s2">, value=</span><span class="se">\n</span><span class="s2">"</span> |
| 1063 | + <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">indented_value_repr</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span> |
| 1064 | + <span class="sa">f</span><span class="s2">")"</span> |
| 1065 | + <span class="p">)</span> |
| 1066 | + <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">is_gradtrackingtensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">):</span> |
| 1067 | + <span class="k">return</span> <span class="p">(</span> |
| 1068 | + <span class="sa">f</span><span class="s2">"GradTrackingTensor(lvl=</span><span class="si">{</span><span class="n">level</span><span class="si">}</span><span class="s2">, value=</span><span class="se">\n</span><span class="s2">"</span> <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">indented_value_repr</span><span class="si">}</span><span class="se">\n</span><span class="s2">"</span> <span class="sa">f</span><span class="s2">")"</span> |
| 1069 | + <span class="p">)</span> |
| 1070 | + <span class="k">if</span> <span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_functorch</span><span class="o">.</span><span class="n">is_functionaltensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">):</span> |
| 1071 | + <span class="k">return</span> <span class="sa">f</span><span class="s2">"FunctionalTensor(lvl=</span><span class="si">{</span><span class="n">level</span><span class="si">}</span><span class="s2">, value=</span><span class="se">\\\n</span><span class="si">{</span><span class="n">value_repr</span><span class="si">}</span><span class="s2">)"</span> |
| 1072 | + |
| 1073 | + <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">"We don't know how to print this, please file us an issue"</span><span class="p">)</span> |
| 1074 | + |
| 1075 | + |
1031 | 1076 | <span class="k">def</span> <span class="nf">_str</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="p">,</span> <span class="n">tensor_contents</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
|
1032 | 1077 | <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">no_grad</span><span class="p">():</span>
|
1033 | 1078 | <span class="k">return</span> <span class="n">_str_intern</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">tensor_contents</span><span class="o">=</span><span class="n">tensor_contents</span><span class="p">)</span>
|
|
0 commit comments