[Torch] Support Python list, more realistic recurrent networks

[masahi]

This PR builds on the control flow support added in #4964 and aims to support more realistic recurrent networks than the simple one in #4964. Specifically the goal is to enable translating LSTM models in the PyTorch repo https://github.com/pytorch/pytorch/tree/master/benchmarks/fastrnns described in their blog post.

Translating these models requires taking care of dynamic lists and tensor shape. I added necessary support in the Torch frontend using prelude List ADT, static Tensor array #5103, and Any. Previously we can only translate "tensors in, tensors out" type of models, but now we can ingest more complex inputs such as list of tuple of tensors.

See the new test cases for the kinds of models we can support now. I added some variants of LSTMs:

• LSTM with layer normalization
• Bidirectional
• Stacked
• Stacked and Bidirectional

The result of translating three layer stacked bidirectional LSTMs is dumped here. Even though this model has three loop nest,

for layer in num_layers:
    for dir in ["forward", "backward"]:
        seq_len = input.size(0)
        for i in seq_len:
              ...

the two outer loops are unrolled by Torchscript. So in the Relay IR dump, there are 3 (layers) x 2 (direction) while loops.

please review @kevinthesun @zhiics @MarisaKirisame @icemelon9 @jwfromm @wweic @alexwong
cc @tqchen @jroesch @ajtulloch @junrushao1994

[MarisaKirisame]

LGTM. Is torchscript list immutable or mutable like python's list?

[masahi]

LGTM. Is torchscript list immutable or mutable like python's list?

Yes it is mutable. List append is mapped to aten::append in Torchscript and it is entirely side-effecting operation. See below for a simple module that only does list append and its Torchscript representation.

Even though variables that are updated in the loop are supposed to be passed to prim::Loop op to become loop variables, this does not apply to side effecting operations like list append. Translating this module to Relay is complicated because also in Relay loop variables are the only ones that can be updated between iteration. If we naively translate it, the list outputs.1 below appears free in the loop and can not be updated.

class ListAppend(nn.Module):
    def forward(self, input):
        # type: (Tensor) -> List[Tensor]
        outputs = []
        for i in range(input.size(0)):
            outputs.append(input)
        return outputs

graph(%self : __torch__.ListAppend,
      %input.1 : Tensor):
  %8 : bool = prim::Constant[value=1]() # rnn_test.py:142:8
  %4 : int = prim::Constant[value=0]() # rnn_test.py:142:34
  %outputs.1 : Tensor[] = prim::ListConstruct()
  %5 : int = aten::size(%input.1, %4) # rnn_test.py:142:23
   = prim::Loop(%5, %8) # rnn_test.py:142:8
    block0(%i : int):
      %12 : Tensor[] = aten::append(%outputs.1, %input.1) # rnn_test.py:143:12
      -> (%8)
  return (%outputs.1)

To workaround the difficulty of list append, I use list concat to append one element at the tail of a list. The original LSTM models in Pytorch repo do not use list append either and use concat instead, probably for the same reason.

[MarisaKirisame]

From an outsider, it seems like the more principled approach is to translate list to Reference of List. We could then write passses to remove Reference if possible.

[kevinthesun]

Overall lgtm. Some minor comments regarding to registering static tensor array ops.

[masahi]

@kevinthesun Thanks for the review! Please have a look at the last commit.

[kevinthesun]

LGTM

[kevinthesun]

Thanks @masahi @MarisaKirisame