<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class MyModule(nn.Module)
    number = Optional[int]

    def __init__(self, number=None):
        self.number = number

    def forward(self):
        if self.number is not None:
            self.number += 1

my_module = torch.jit.script(MyModule())"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">MyModule</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>)
    <span class="pl-s1">number</span> <span class="pl-c1">=</span> <span class="pl-v">Optional</span>[<span class="pl-s1">int</span>]

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">number</span><span class="pl-c1">=</span><span class="pl-c1">None</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">number</span> <span class="pl-c1">=</span> <span class="pl-s1">number</span>

    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>):
        <span class="pl-k">if</span> <span class="pl-s1">self</span>.<span class="pl-s1">number</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
            <span class="pl-s1">self</span>.<span class="pl-s1">number</span> <span class="pl-c1">+=</span> <span class="pl-c1">1</span>

<span class="pl-s1">my_module</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-en">script</span>(<span class="pl-v">MyModule</span>())</pre></div>
<p dir="auto">Always have to work around it.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class MyModule(nn.Module)
    number = Optional[int]

    def __init__(self, number=None):
        self.number = number

    def forward(self):
        number = self.number
        if number is not None:
           number += 1

my_module = torch.jit.script(MyModule())"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">MyModule</span>(<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>)
    <span class="pl-s1">number</span> <span class="pl-c1">=</span> <span class="pl-v">Optional</span>[<span class="pl-s1">int</span>]

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">number</span><span class="pl-c1">=</span><span class="pl-c1">None</span>):
        <span class="pl-s1">self</span>.<span class="pl-s1">number</span> <span class="pl-c1">=</span> <span class="pl-s1">number</span>

    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>):
        <span class="pl-s1">number</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">number</span>
        <span class="pl-k">if</span> <span class="pl-s1">number</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
           <span class="pl-s1">number</span> <span class="pl-c1">+=</span> <span class="pl-c1">1</span>

<span class="pl-s1">my_module</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-en">script</span>(<span class="pl-v">MyModule</span>())</pre></div>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/suo/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/suo">@suo</a></p>
<p dir="auto">Right now we only do type refinement on <code class="notranslate">TK_VAR</code>s in the compiler, we should expand this to be more general, these examples below don't currently work:</p>
<p dir="auto">Refinement on an attribute:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class M(torch.jit.ScriptModule):
    def __init__(self):
        super(M, self).__init__()
        self.id = jit.Attribute(2, Optional[int])

    @jit.script_method
    def forward(self):
        a = 2
        if self.id is not None:
            a += self.id

m = M()"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">M</span>(<span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-v">ScriptModule</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>):
        <span class="pl-en">super</span>(<span class="pl-v">M</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">id</span> <span class="pl-c1">=</span> <span class="pl-s1">jit</span>.<span class="pl-v">Attribute</span>(<span class="pl-c1">2</span>, <span class="pl-v">Optional</span>[<span class="pl-s1">int</span>])

    <span class="pl-en">@<span class="pl-s1">jit</span>.<span class="pl-s1">script_method</span></span>
    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>):
        <span class="pl-s1">a</span> <span class="pl-c1">=</span> <span class="pl-c1">2</span>
        <span class="pl-k">if</span> <span class="pl-s1">self</span>.<span class="pl-s1">id</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
            <span class="pl-s1">a</span> <span class="pl-c1">+=</span> <span class="pl-s1">self</span>.<span class="pl-s1">id</span>

<span class="pl-s1">m</span> <span class="pl-c1">=</span> <span class="pl-v">M</span>()</pre></div>
<p dir="auto">Refinement on a constant:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="class M(torch.jit.ScriptModule):
    __constants__ = ['x']

    def __init__(self, x):
        super(M, self).__init__()
        self.x = x

    @torch.jit.script_method
    def forward(self, a):
        attr = self.x
        if attr is not None:
            return a + attr
        else:
            return a

M(None)"><pre class="notranslate"><span class="pl-k">class</span> <span class="pl-v">M</span>(<span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-v">ScriptModule</span>):
    <span class="pl-s1">__constants__</span> <span class="pl-c1">=</span> [<span class="pl-s">'x'</span>]

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-en">super</span>(<span class="pl-v">M</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span>

    <span class="pl-en">@<span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-s1">script_method</span></span>
    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">a</span>):
        <span class="pl-s1">attr</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">x</span>
        <span class="pl-k">if</span> <span class="pl-s1">attr</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>:
            <span class="pl-k">return</span> <span class="pl-s1">a</span> <span class="pl-c1">+</span> <span class="pl-s1">attr</span>
        <span class="pl-k">else</span>:
            <span class="pl-k">return</span> <span class="pl-s1">a</span>

<span class="pl-v">M</span>(<span class="pl-c1">None</span>)</pre></div>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gmagogsfm/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gmagogsfm">@gmagogsfm</a></p>