Support aten::lstm, aten::gru and aten::rnn_{relu,tanh} operations (#21579)

* [PT FE] Initial LSTM impl

* Support aten::lstm operation

* Support aten::gru and aten::rnn_{relu,tanh}

* Apply suggestions from code review

* Fix build

* Simplify implementation

* More simplification

* Update src/frontends/pytorch/src/op/lstm.cpp

* Simplification

* Use default LSTM helpers

---------

Co-authored-by: Roboreptile <piotr.krzeminski@intel.com>

tests/layer_tests/pytorch_tests/test_lstm.py

@@ -0,0 +1,140 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class aten_lstm(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first):
+        torch.nn.Module.__init__(self)
+        self.lstm = torch.nn.LSTM(input_size,
+                                  hidden_size,
+                                  num_layers,
+                                  has_bias,
+                                  batch_first,
+                                  bidirectional=bidirectional)
+
+    def forward(self, input_tensor, h0, c0):
+        return self.lstm(input_tensor, (h0, c0))
+
+
+class aten_gru(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first):
+        torch.nn.Module.__init__(self)
+        self.gru = torch.nn.GRU(input_size,
+                                hidden_size,
+                                num_layers,
+                                has_bias,
+                                batch_first,
+                                bidirectional=bidirectional)
+
+    def forward(self, input_tensor, h0):
+        return self.gru(input_tensor, h0)
+
+
+class aten_rnn(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, nonlinearity):
+        torch.nn.Module.__init__(self)
+        self.rnn = torch.nn.RNN(input_size,
+                                hidden_size,
+                                num_layers,
+                                nonlinearity=nonlinearity,
+                                bias=has_bias,
+                                batch_first=batch_first,
+                                bidirectional=bidirectional)
+
+    def forward(self, input_tensor, h0):
+        return self.rnn(input_tensor, h0)
+
+
+class TestLSTM(PytorchLayerTest):
+    def _prepare_input(self):
+        n = self.num_layers
+        if self.bidirectional:
+            n *= 2
+        if self.batch_first:
+            input = np.random.randn(3, 5, self.input_size).astype(np.float32)
+        else:
+            input = np.random.randn(5, 3, self.input_size).astype(np.float32)
+        h0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
+        c0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
+        return (input, h0, c0)
+
+    @pytest.mark.parametrize("input_size,hidden_size", [(10, 20),])
+    @pytest.mark.parametrize("num_layers", [1, 2, 7])
+    @pytest.mark.parametrize("has_bias", [True, False])
+    @pytest.mark.parametrize("bidirectional", [True, False])
+    @pytest.mark.parametrize("batch_first", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_lstm(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, ie_device, precision, ir_version):
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.batch_first = batch_first
+        self._test(aten_lstm(input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first), None, "aten::lstm",
+                   ie_device, precision, ir_version, trace_model=True)
+
+
+class TestGRU(PytorchLayerTest):
+    def _prepare_input(self):
+        n = self.num_layers
+        if self.bidirectional:
+            n *= 2
+        if self.batch_first:
+            input = np.random.randn(3, 5, self.input_size).astype(np.float32)
+        else:
+            input = np.random.randn(5, 3, self.input_size).astype(np.float32)
+        h0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
+        return (input, h0)
+
+    @pytest.mark.parametrize("input_size,hidden_size", [(10, 20),])
+    @pytest.mark.parametrize("num_layers", [1, 2, 7])
+    @pytest.mark.parametrize("has_bias", [True, False])
+    @pytest.mark.parametrize("bidirectional", [True, False])
+    @pytest.mark.parametrize("batch_first", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_gru(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, ie_device, precision, ir_version):
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.batch_first = batch_first
+        self._test(aten_gru(input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first), None, "aten::gru",
+                   ie_device, precision, ir_version, trace_model=True)
+
+
+class TestRNN(PytorchLayerTest):
+    def _prepare_input(self):
+        n = self.num_layers
+        if self.bidirectional:
+            n *= 2
+        if self.batch_first:
+            input = np.random.randn(3, 5, self.input_size).astype(np.float32)
+        else:
+            input = np.random.randn(5, 3, self.input_size).astype(np.float32)
+        h0 = np.random.randn(n, 3, self.hidden_size).astype(np.float32)
+        return (input, h0)
+
+    @pytest.mark.parametrize("input_size,hidden_size", [(10, 20),])
+    @pytest.mark.parametrize("num_layers", [1, 2, 7])
+    @pytest.mark.parametrize("has_bias", [True, False])
+    @pytest.mark.parametrize("bidirectional", [True, False])
+    @pytest.mark.parametrize("batch_first", [True, False])
+    @pytest.mark.parametrize("nonlinearity", ["tanh", "relu"])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_rnn(self, input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, nonlinearity, ie_device, precision, ir_version):
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.batch_first = batch_first
+        self._test(aten_rnn(input_size, hidden_size, num_layers, has_bias, bidirectional, batch_first, nonlinearity), None, f"aten::rnn_{nonlinearity}",
+                   ie_device, precision, ir_version, trace_model=True)

tests/model_hub_tests/torch_tests/hf_transformers_models

@@ -67,7 +67,7 @@ facebook/convnextv2-tiny-22k-384,convnextv2
 facebook/detr-resnet-50,detr
 facebook/dinov2-base,dinov2,xfail,Tracing error: Please check correctness of provided example_input (but eval was correct)
 facebook/dpr-question_encoder-single-nq-base,dpr
-facebook/encodec_24khz,encodec,xfail,Unsupported op aten::lstm
+facebook/encodec_24khz,encodec
 facebook/esm2_t6_8M_UR50D,esm
 facebook/flava-full,flava,xfail,Tracing problem
 facebook/flava-image-codebook,flava_image_codebook,skip,Load problem

tests/model_hub_tests/torch_tests/test_timm.py

@@ -68,7 +68,8 @@ class TestTimmConvertModel(TestTorchConvertModel):
     @pytest.mark.parametrize("name", ["mobilevitv2_050.cvnets_in1k",
                                       "poolformerv2_s12.sail_in1k",
                                       "vit_base_patch8_224.augreg_in21k",
-                                      "beit_base_patch16_224.in22k_ft_in22k"])
+                                      "beit_base_patch16_224.in22k_ft_in22k",
+                                      "sequencer2d_l.in1k"])
     @pytest.mark.precommit
     def test_convert_model_precommit(self, name, ie_device):
         self.run(name, None, ie_device)

tests/model_hub_tests/torch_tests/timm_models

@@ -385,7 +385,7 @@ selecsls60.in1k,None
 selecsls60b.in1k,None
 semnasnet_075.rmsp_in1k,None
 senet154.gluon_in1k,None
-sequencer2d_l.in1k,None,xfail,Unsupported aten::lstm
+sequencer2d_l.in1k,None
 seresnet152d.ra2_in1k,None
 seresnet33ts.ra2_in1k,None
 seresnet50.a1_in1k,None

tests/model_hub_tests/torch_tests/torchbench_models

@@ -10,7 +10,7 @@ basic_gnn_gin,None,xfail,Unsupported op aten::scatter_add_
 basic_gnn_sage,None,xfail,Unsupported op aten::scatter_add_
 #cm3leon_generate,None,skip,No install.py is found
 dcgan,None
-demucs,None,xfail,Unsupported op aten::lstm
+demucs,None
 #densenet121,None - Already tested by torchvision tests
 #detectron2_fasterrcnn_r_101_c4,None - Already tested by det2 tests
 #detectron2_fasterrcnn_r_101_dc5,None - Already tested by det2 tests
@@ -93,7 +93,7 @@ tacotron2,None,skip,Can't be loaded without CUDA
 #timm_vision_transformer_large,None - Already tested by timm tests
 #timm_vovnet,None - Already tested by timm tests
 torch_multimodal_clip,None,skip,Can't be traced
-tts_angular,None,xfail,Unsupported op aten::lstm
+tts_angular,None
 #vgg16,None - Already tested by torchvision tests
 #vision_maskrcnn,None,skip,Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
 yolov3,None