Is there support of recursive custom classes like Tree in Tree-LSTM?

username: FereX98
cooked: <p>Hi there!<br>
I am currently trying to make JIT optimizations work on the source code of Tree-LSTM model. The Tree class in the model is a crucial part of it, so I need to make it a custom class type so it can be used by the core methods of the model. That’s when I find a problem:</p>
<pre><code class="lang-auto">import torch
from typing import List

@torch.jit.script
class Tree(object):
    def __init__(self):
        self.parent = None
        self.num_children = 0
        self.children = torch.jit.annotate(List[Tree], [])

    # further definitions omitted
</code></pre>
<p>When I try to run the code, here is the error:</p>
<pre><code class="lang-auto">RuntimeError: 
Unknown type name Tree:
def __init__(self):
    self.parent = None
    self.num_children = 0
    self.children = torch.jit.annotate(List[Tree], [])
                                            ~~~~ &lt;--- HERE
</code></pre>
<p>So the question is: Tree is basically a recursive structure, the children of a tree node is a list of tree nodes. Therefore for the children variable of a Tree class, I need to define an empty list of Tree class type. But since the Tree class type definition is still halfway, the interpreter cannot recognize Tree type.</p>
<p>I am wondering is there any method that I can solve this problem, or there is actually no support for custom classes like the Tree above in current version of PyTorch and I should try other ways.</p>
<p>It would be so nice if someone can give me a hand. Thanks a lot!</p>

username: Michael_Suo
cooked: <p>I believe this is fixed in master (your class compiles for me). Could you try your code on our nightly release and see if it works?</p>

username: FereX98
cooked: <p>The code successfully runs after I’ve switched to nightly release, cool! Thank you!</p>
<p>But there is bad news that new errors arise when the interpreter tries to compile other methods. For example:</p>
<pre><code class="lang-auto">@torch.jit.script
class Tree(object):
    def __init__(self):
        self.parent = torch.jit.annotate(Optional[Tree], None)

    def add_child(self, child):
        # type: (Tree) -&gt; None
        child.parent = torch.jit.annotate(Optional[Tree], self)
</code></pre>
<p>When run to the last line, an error arises:</p>
<pre><code class="lang-auto">RuntimeError: 
expected an expression of type Optional[__torch__.Tree] but found __torch__.Tree:
at /Users/fere/repos/treelstm.pytorch/treelstm/tree.py:25:43
    def add_child(self, child):
        # type: (Tree) -&gt; None
        child.parent = torch.jit.annotate(Optional[Tree], self)
                                          ~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
</code></pre>
<p>The interpreter regards <code>self</code> as <code>Tree</code> instead of <code>Optional[Tree]</code>, which makes sense but I thought that <code>torch.jit.annotate</code> should be able to do some casting work. The reason why I think so is that in code above, the interpreter takes <code>torch.jit.annotate(Optional[Tree], None)</code> as an <code>Optional[Tree]</code> type value instead of <code>None</code>. I’ve gone through the TorchScript documentation and done some searching, but the only method I’ve found that is able to specify types for a value other than function parameters and return values is by using <code>torch.jit.annotate</code>. Also, I haven’t find any ways which can cast a <code>T</code> typed value to an <code>Optional[T]</code> value. So I am curious that is there anything I miss?</p>
<p>Another problem is below:</p>
<pre><code class="lang-auto">@torch.jit.script
class Tree(object):
    def __init__(self):
        self.num_children = torch.jit.annotate(int, 0)

    def add_child(self):
        self.num_children += 1
</code></pre>
<p>Error:</p>
<pre><code class="lang-auto">RuntimeError: 
left-hand side of augmented assignment to module parameters/buffers can only be tensor types:
at /Users/fere/repos/treelstm.pytorch/treelstm/tree.py:15:9
    def add_child(self):
        self.num_children += 1
        ~~~~~~~~ &lt;--- HERE
</code></pre>
<p>Does the error message means that only Tensor-typed values can be updated in a custom TorchScript class?</p>
<p>Thankful and grateful for your help!</p>

username: Michael_Suo
cooked: <p>First of all—thank you for the feedback! Support for classes is in its early days and reports from intrepid users are super valuable.</p>
<p>To summarize, it seems that you are running into two problems:</p>
<ol>
<li>
<p><code>torch.jit.annotate()</code> is not doing implicit type promotion from <code>T</code> to <code>Optional[T]</code> for class types.<br>
This is a bug on our side, and we’ll look into fixing it. Can you file a Github task and we can track it there?</p>
</li>
<li>
<p>Augmented assignment doesn’t work on non-tensors. This is a known issue, and we have a fix coming this week for it. In the meantime, just re-assigning is an easy workaround:</p>
</li>
</ol>
<pre><code class="lang-auto">def add_child(self):
    self.num_children = self.num_children + 1
</code></pre>

username: Michael_Suo
cooked: <p>actually no need to file an issue, <a href="https://github.com/pytorch/pytorch/pull/21593" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/21593</a> fixes <span class="hashtag">#1</span>.</p>

username: Michael_Suo
cooked: <p>Update:<br>
We do not actually currently support recursive class definitions. I didn’t remember when I first replied. <a href="https://github.com/pytorch/pytorch/pull/21842" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/21842</a> improves the error message in this case.</p>
<p>We will likely support it soon (before the next release) but for now it will not work.</p>

username: FereX98
cooked: <p>Thanks a lot for the update! Guess now I need to wait or look for a temporal workaround. <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=9" title=":grinning:" class="emoji" alt=":grinning:"></p>