[Feature Request][Relax][Torch] Add support for nn.LSTM (lstm.input) in from_exported_program #18340

import torch
import torch.nn as nn
from torch.export import export as torch_export
from tvm.relax.frontend.torch import from_exported_program

class M(nn.Module):
    def __init__(self, input_size=8, hidden_size=16, num_layers=1, bidirectional=False):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=bidirectional,
        )

    def forward(self, x):
        # Only return the output sequence; drop (h_n, c_n)
        y, _ = self.lstm(x)
        return y

def main():
    torch.manual_seed(0)
    m = M().eval()

    # Minimal input: B=2, T=4, C=8
    x = torch.randn(2, 4, 8, dtype=torch.float32)

    # 1) Sanity check in eager
    with torch.inference_mode():
        y = m(x)
    print("PyTorch eager OK, y.shape =", tuple(y.shape))

    # 2) Export to ExportedProgram
    ep = torch_export(m, (x,))
    print("ExportedProgram created.")

    # 3) Import into TVM Relax — triggers unsupported function type
    _ = from_exported_program(ep)

if __name__ == "__main__":
    main()