[JIT] [Mobile] Is isinstance() supposed to work with TorchScript

username: zetyquickly
cooked: <p>Hello there,</p>
<p>I’ve faced similar issues during deploying scripted model on Android.</p>
<ol>
<li>The first question is “How to check whole grap or code of scripted model?”.<br>
<code>ScriptedModule.code</code> provides only upper level code, but inside stacktrace on mobile I see more information about how the code looks like.</li>
<li>Second is about isinstance(). My model has plenty of <code>isinstance(tensor, torch.Tensor)</code> or such which are converted to in <code>tensor1 = unchecked_cast(Tensor, tensor)</code> script code. Here comes an error.</li>
<li>Why there’s no error while exporting/scripting of module?</li>
</ol>

username: driazati
cooked: <p><code>isinstance</code> is supported but its result is static. It is useful for <code>Module</code>s that have different attribute types passed in, e.g.</p>
<pre><code class="lang-python">class M(torch.nn.Module):
    def __init__(self, x):
        super().__init__()
        self.x = x

    def forward(self):
        if isinstance(self.x, List[str]):
            return self.x[2]
        else:
            return self.x + 2

print(torch.jit.script(M(['bye'])).graph)
print(torch.jit.script(M(2)).graph)
</code></pre>
<p>The compiler is able to see the <code>isinstance</code> check and evaluate it at compile time and remove the unused branch. The graphs show this:</p>
<pre><code class="lang-auto">graph(%self : ClassType&lt;M&gt;):
  %3 : str = prim::Constant[value="hi"]() # ../test.py:22:19
  return (%3)

graph(%self : ClassType&lt;M&gt;):
  %4 : int = prim::Constant[value=2]() # ../test.py:24:28
  %3 : int = prim::GetAttr[name="x"](%self)
  %5 : int = aten::add(%3, %4) # ../test.py:24:19
  return (%5)
</code></pre>

username: driazati
cooked: <p>As for 1), we recently changed the behavior so that functions in <code>.code</code> and <code>.graph</code> appear as function calls (previously we were inlining the function bodies). So we’re still missing the functionality to show the called functions. For now you can re-enable inlining to see the entire graph:</p>
<pre><code class="lang-python">def other_fn(x):
    return x + 10

# Change the inlining mode before you compile
torch._C._jit_set_inline_everything_mode(True)

@torch.jit.script
def fn(x):
    return other_fn(x)

print(fn.code)
print(fn.graph)
</code></pre>
<p>You can track this bug in <a href="https://github.com/pytorch/pytorch/issues/29750" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/29750</a>.</p>

username: driazati
cooked: <aside class="quote no-group" data-username="driazati" data-post="3" data-topic="60845">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/driazati/48/15393_2.png" class="avatar"> driazati:</div>
<blockquote>
<p>def other_fn(x): return x + 10 # Change the inlining mode before you compile torch._C._jit_set_inline_everything_mode(True) <span class="mention">@torch.jit.script</span> def fn(x): return other_fn(x) print(fn.code) print(fn.graph)</p>
</blockquote>
</aside>
<p>You can also print out a model directly with</p>
<pre data-code-wrap="python"><code class="lang-python">class M(nn.Module):
    def other_fn(self, x):
        return x + 10

    def forward(self, x):
        return self.other_fn(x)

m = torch.jit.script(M())
print(m._c.dump())
</code></pre>

username: zetyquickly
cooked: <p>This one is unclear to me.</p>
<p>How could we annotate input for <code>__init__()</code> and output of <code>forward()</code> functions? The <code>Union[int, List[str]]</code> typing is unsupported</p>

username: zetyquickly
cooked: <aside class="quote no-group" data-username="driazati" data-post="3" data-topic="60845">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/driazati/48/15393_2.png" class="avatar"> driazati:</div>
<blockquote>
<p>torch._C._jit_set_inline_everything_mode(True)</p>
</blockquote>
</aside>
<p>Thanks a lot! That clarified the way how to see the whole code</p>

username: zetyquickly
cooked: <p>This returns <code>None</code> but anyway I haven’t looked for that functionality <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>

username: driazati
cooked: <aside class="quote no-group" data-username="zetyquickly" data-post="5" data-topic="60845" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zetyquickly/48/22489_2.png" class="avatar"> zetyquickly:</div>
<blockquote>
<p>This one is unclear to me.</p>
<p>How could we annotate input for <code>__init__()</code> and output of <code>forward()</code> functions? The <code>Union[int, List[str]]</code> typing is unsupported</p>
</blockquote>
</aside>
<p><code>__init__</code> on <code>nn.Module</code>s runs in Python (<code>torch.jit.script</code> only sees the module after it has been initialized), so you can annotate that with whatever Python annotations you want (but they won’t be enforced by the compiler at all, for that you should use something like mypy).</p>
<p>For methods that are compiled (e.g. <code>forward</code> and anything it calls), the return types can be deduced from the code. If you want to explicitly write it out, you can use any of the type annotations listed <a href="https://pytorch.org/docs/master/jit.html#supported-type">here</a>.</p>
<p><code>Union</code>s aren’t supported so they won’t work in TorchScript. As a workaround you could do something like <code>Tuple[Optional[int], Optional[List[str]]]</code>.</p>