Unpack tuple/list input

[csukuangfj]

For some networks, e.g., RNN/LSTM, we need to provide some previous states as input.

Example code

#!/usr/bin/env python3

import torch
import torch.nn as nn
from typing import List


class Foo(nn.Module):
    def forward(self, x, states: List[torch.Tensor]):
        return x, states[0] + states[1]


def main():
    f = Foo()

    x = torch.rand(2, 3, 4)
    y = [x, x]
    m = torch.jit.trace(f, (x, y))
    print(m.graph)
    m.save("m.pt")


if __name__ == "__main__":
    torch.manual_seed(20221128)
    main()

[nihui]

we need level0 pass otherwise shape inference will fail

[csukuangfj]

we need level0 pass otherwise shape inference will fail

Do you mean we should move it to pass_level0 ?

[nihui]

we need level0 pass otherwise shape inference will fail

Do you mean we should move it to pass_level0 ?

yeah, implement flatten inputs on torch graph, and apply the pass before shape inference

reference code
https://github.com/pytorch/pytorch/blob/3120054c151d37fdf43963dbb60ab420908f48cf/torch/csrc/jit/passes/lower_tuples.cpp#L185

[nihui]

move to #4498