Support recursive data type in TorchScript

[yf225]

🚀 Feature

It would be really awesome to support recursive data type in TorchScript. For example:

import torch
from typing import Dict

class TypedDataDict(object):
  def __init__(self):
    self.str_to_dict: Dict[str, 'TypedDataDict'] = {}

  def set_str_to_dict(self, value: Dict[str, 'TypedDataDict']):
    self.str_to_dict = value


class TestModule(torch.nn.Module):
  def __init__(self):
    super().__init__()

  def forward(self, input):
    return TypedDataDict().set_str_to_dict({"123": TypedDataDict()})


m = TestModule()
m_scripted = torch.jit.script(m)
m_scripted(torch.tensor(1.))

'''
Currently throws:

RuntimeError: 
Assignment to attribute 'str_to_dict' cannot be of a type that contains class '__torch__.TypedDataDict'.
Classes that recursively contain instances of themselves are not yet supported:
  File "test_yf225.py", line 6
  def __init__(self):
    self.str_to_dict: Dict[str, 'TypedDataDict'] = {}
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
'TypedDataDict.__init__' is being compiled since it was called from '__torch__.TypedDataDict'
  File "test_yf225.py", line 17
  def forward(self, input):
    return TypedDataDict().set_str_to_dict({"123": TypedDataDict()})
           ~~~~~~~~~~~~~ <--- HERE
'__torch__.TypedDataDict' is being compiled since it was called from 'TestModule.forward'
  File "test_yf225.py", line 17
  def forward(self, input):
    return TypedDataDict().set_str_to_dict({"123": TypedDataDict()})
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
'''

Motivation

Many DPER3 modules take input of nested data types like Dict[str, Dict[str, Dict[str, torch.Tensor]]], and it will be really hard to maintain if we had to add support to all variations of those nested data types to TypedDataDict (which is DPER3's typed dictionary object class). But if recursive data type is supported, then we only need Dict[str, TypedDataDict] and Dict[str, torch.Tensor] to cover all nesting possibilities.

Having recursive data type support will greatly speed up the work for moving all PyPer models to 100% TorchScript, demonstrating TorchScript's production readiness for large-scale ranking models.

cc. @wanchaol @suo

cc @ezyang @gchanan @zou3519 @suo @gmagogsfm

[gmagogsfm]

Discussed with @suo offline, we will need to sync with yf225 more to figure out a solution.

[0xJchen]

Discussed with @suo offline, we will need to sync with yf225 more to figure out a solution.

Hi, guys. I wonder if there are any workouts to avoid this problem?

I found recursive data types are common in many algorithms (as mentioned in PEP 484). Take tree search as an example, we have a node class, and it also has children (let's say, a mapping from index to new nodes) self.children: Dict[int, Node]={}. I met the above problem when trying to convert a tree-search algorithm to Torchscript, and I wonder if there are any ways to avoid it?

This is an example:

@torch.jit.script
class Node(object):

    def __init__(self, prior: float):
        self.children: Dict[int, 'Node'] = { }

    def expanded(self) -> bool:
        return len(self.children) > 0

    def expand(self, priors):
        for id, prior in enumerate(priors):
            self.children[id] = Node(prior )

Also, I wonder if I wrote a Cython script for managing the above tree search logic (I can directly call those Cython functions after compilation in a normal python script), can I still re-use them in Torchscript?

cc @yf225 @suo @gmagogsfm