No graph attribute when use torch.jit._overload_method

BobChen: <p>PyTorch version: 1.6.0.dev20200407+cu101</p>
<p>When use torch.jit._overload_method to overload functions, it seems that there is no graph attribute in the output torchscript.</p>
<pre><code class="lang-auto">class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

    @torch.jit._overload_method  # noqa: F811
    def forward(self, x: int) -&gt; int:  # noqa: F811
        pass

    @torch.jit._overload_method  # noqa: F811
    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:  # noqa: F811
        pass

    def forward(self, x):  # noqa: F811
        original = x

        if isinstance(original, int):
            return original + 2
        else:
            return original.sum()

model = MyModule()
s = torch.jit.script(model)
torch._C._jit_pass_inline(s.graph)  # error happens here
</code></pre>
<p>error message:</p>
<pre><code class="lang-auto">torch.nn.modules.module.ModuleAttributeError: 'RecursiveScriptModule' object has no attribute 'graph'
</code></pre>

eellison: <p>Hi, this is is because there is no single “graph”, there are two separate graphs, one for each overload. I could maybe add an api for something like <code>model.forwards.graphs</code>, or <code>model.forward.graph_for_types(1)</code> to get the graph for a specific overload.</p>
<p>Generally, overloads should be usable and work as you expect when they are contained in other modules but you may run into a little bit of difficulty if you are interacting with them as a top level module. They’re still an internal feature but will be cleaned up and released probably in the next release or the one after.</p>

BobChen: <aside class="quote no-group" data-username="eellison" data-post="2" data-topic="76128">
<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/eellison/48/9338_2.png" class="avatar"> eellison:</div>
<blockquote>
<p>Generally, overloads should be usable and work as you expect when they are contained in other modules but you may run into a little bit of difficulty if you are interacting with them as a top level module. They’re still an internal feature but will be cleaned up and released probably in the next release or the one after.</p>
</blockquote>
</aside>
<p>Thank you very much:)<br>
I put overloads into sub modules and find there is one single graph for one kind of overload as you said:</p>
<pre><code class="lang-auto">class MySubModule(nn.Module):
    def __init__(self):
        super().__init__()

    @torch.jit._overload_method  # noqa: F811
    def forward(self, x: int) -&gt; int:  # noqa: F811
        pass

    @torch.jit._overload_method  # noqa: F811
    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:  # noqa: F811
        pass

    def forward(self, x):  # noqa: F811
        original = x

        if isinstance(original, int):
            return original + 2
        else:
            return original.sum()

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.sub_module = MySubModule()

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


model = MyModule()
s = torch.jit.script(model)
torch._C._jit_pass_inline(s.graph)
print(s.graph)
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">graph(%self : __torch__.MyModule,
      %x.1 : Tensor):
      %2 : __torch__.MySubModule = prim::GetAttr[name="sub_module"](%self)
      %14 : None = prim::Constant()
      %15 : Tensor = aten::sum(%x.1, %14) 
  return (%15)
</code></pre>
<p>Since default type is torch.Tensor, the corresponding graph is obtained. So is there any ways to obtain one single torchscript for all overloads for now?</p>

eellison: <p>The MyModule you’re using only takes a Tensor. If you want to access the int graph you could try using a different module with the submodule that takes in an int.</p>

BobChen: <p>Got it, thank you very much:)</p>

Padarn_Wilson: <p>Sorry to dredge this one up but I had exactly the same problem, I was hoping I could ask <a class="mention" href="/u/eellison">@eellison</a> how this works exactly?</p>
<blockquote>
<p>Generally, overloads should be usable and work as you expect when they are contained in other modules</p>
</blockquote>
<p>I get the idea: Once you call forward, then the torchscript can inspect types and choose the right graph? But is it possible to get a reference for some documentation to fully understand how this works (or to the code that does this)? I don’t quite understand why forward is treated differently - when this error occurs I can still call the other functions which have been overloaded.</p>

Elias_Ellison: <p>Hi <a class="mention" href="/u/padarn_wilson">@Padarn_Wilson</a>, the reason there aren’t good docs for this is because it wasn’t fully finished. there are still a few rough edges as you are encountering.</p>
<blockquote>
<p>I don’t quite understand why forward is treated differently - when this error occurs I can still call the other functions which have been overloaded.</p>
</blockquote>
<p>Hmm, what do you mean exactly ? Not sure I follow.</p>

Padarn_Wilson: <p>Hey <a class="mention" href="/u/elias_ellison">@Elias_Ellison</a> thanks for the calrification.</p>
<aside class="quote no-group" data-username="Elias_Ellison" data-post="7" data-topic="76128">
<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/elias_ellison/48/39583_2.png" class="avatar"> Elias_Ellison:</div>
<blockquote>
<p>Hmm, what do you mean exactly ? Not sure I follow.</p>
</blockquote>
</aside>
<p>From the discussion above I understood that when <code>.forward</code> is invoked it will also do the job of inferring which overloaded functions to call based on the types present… but rereading it I realise I had my understanding a bit confused. I think it is clear now.</p>