Should adding @torch.jit.export annotation on a single method make it export all module properties?

pkubik: <p>I run into this interesting issue when trying to export a simple <code>nn.Module</code>. Here is the code:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class MyMod(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.x = 1
        self.y = 2

    def forward(self, x):
        return 2 * x

    # @torch.jit.export  # &lt;-- UNCOMMENT HERE
    def get_x(self):
        return self.x


mod = MyMod()
traced = torch.jit.trace(mod, [torch.tensor([1.0])])

print(traced.get_x())
print(traced.x)
print(traced.y)
</code></pre>
<p>At the current form the code fails on all 3 print statements - AS EXPECTED.</p>
<p>Now when I uncomment the <code>torch.jit.export</code> annotation all 3 statements pass.</p>
<ul>
<li>I understand <code>get_x()</code>, because I explicitly exported it.</li>
<li>I kind of get <code>x</code>, because it’s a dependency for <code>get_x()</code></li>
<li>I totally don’t get <code>y</code>, because it’s just an independent non-tensor property.</li>
</ul>
<p>Tested on torch==2.0.1.</p>
<p>Most of the time I wouldn’t really care about some extra properties exported alongside my model.<br>
The problem occurs when I try to use some external libraries like HuggingFace <code>transformers</code>.<br>
For example, when I add <code>ViTForImageClassification</code> as a submodule of <code>MyMod</code> all it’s properties<br>
are going to be recursively traced and this leads to annotation errors like the one below:</p>
<pre><code class="lang-auto">Unknown type name 'nn.Module':
  File "blah/site-packages/transformers/modeling_utils.py", line 1197
    @property
    def base_model(self) -&gt; nn.Module:
</code></pre>
<p>This is just some innocent annotated property that gets cought in the crossfire.</p>
<hr>
<p>Any idea if this is an expected behaior. Should I report a bug? I already found an ugly workaround, but this feels like something that should have a simpler solution.</p>
<details>
<summary>
Workaround for anyone interested</summary>
<p>Just create a new <code>MyExports</code> module and move the exports there. Since the behavior is recursive it will be limited to the <code>MyExports</code> module and leave the ViT model alone. You can then access them like <code>mod.exports.get_x()</code>.</p>
</details>

pkubik: <p>I just added <code>print(type(traced))</code> to the code and saw that:</p>
<ul>
<li>without export decorator it’s <code>&lt;class 'torch.jit._trace.TopLevelTracedModule'&gt;</code></li>
<li>with decorator it’s <code>&lt;class 'torch.jit._script.RecursiveScriptModule'&gt;</code></li>
</ul>
<p>So even the result type differes. Is it the case that <code>torch.jit.trace</code> silently enters scripting mode when it sees any export? It’s still weird, because it doesn’t behave exactly like <code>torch.jit.script</code>, but it gets close.</p>
<p>I also noticed that I can get what I want by calling <code>torch.jit.trace_module</code> explicitly like this:</p>
<pre><code class="lang-auto">traced = torch.jit.trace_module(mod, {"forward": [torch.tensor([1.0])], "get_x": []})
</code></pre>
<p>Is there a decorator that I could put on <code>get_x</code> to achieve the same result with <code>torch.jit.trace</code>?</p>

pkubik: <p>I found the piece of code responsible for this behavior:</p>
<pre><code class="lang-auto">def make_module(mod, _module_class, _compilation_unit):
    if isinstance(mod, ScriptModule):
        return mod
    elif torch._jit_internal.module_has_exports(mod):

        infer_methods_stubs_fn = torch.jit._recursive.make_stubs_from_exported_methods
        return torch.jit._recursive.create_script_module(
            mod,
            infer_methods_stubs_fn,
            share_types=False,
            is_tracing=True
        )
    else:
        if _module_class is None:
            _module_class = TopLevelTracedModule
        return _module_class(mod, _compilation_unit=_compilation_unit)
</code></pre>
<p>This is called when tracing a module and it:</p>
<ul>
<li>checks whether the module contains any export decoracors</li>
<li>if yes, then calls <code>create_script_module</code></li>
<li>if not, then it does the regular tracing?</li>
</ul>
<p>Not sure what’s the difference between regular <code>torch.jit.script</code> and the way it is done here, but there is this <code>is_tracing=True</code> flag, which probably makes the difference.</p>
<p>So it seems like this is explicitly coded behavior, but I’m not sure if it’s really desired. If someone explicitly uses tracing then there is probably a reason to that. It would be probably safer to throw an error that you can’t use exports when tracing. I don’t know.</p>