quantization: make lstm quantizable when has state input (#1473)

Author: XiaobingSuper

intel_extension_for_pytorch/quantization/_recipe.py

@@ -66,7 +66,7 @@ def _default_recipe_init(nodes):
 
             # For LSTM, if it's input is a PackedSequence, we don't support ot now.
             # TODO: support PackedSequence input for quantization LSTM.
-            if node.type in rnn_ops and len(node.input_tensor_infos) > 2:
+            if node.type in rnn_ops and len(node.input_tensor_infos) > 2 and node.input_tensor_infos[1].orig_dtype == torch.int64:
                 for idx, tensor_info in enumerate(node.input_tensor_infos):
                     if tensor_info is not None:
                         tensor_info.inf_dtype = tensor_info.orig_dtype

intel_extension_for_pytorch/quantization/_utils.py

@@ -435,7 +435,7 @@ def _reset_post_node_input_infos(node):
                     _reset_post_node_input_infos(node)
             else:
                 # TODO: enable PackedSequence input for LSTM.
-                if not (node.type in [nn.LSTM] and len(node.input_tensor_infos) > 2):
+                if not (node.type in [nn.LSTM] and len(node.input_tensor_infos) > 2 and node.input_tensor_infos[1].orig_dtype == torch.int64):
                     if node.input_tensor_force_inf_dtype[0] in [torch.qint8, torch.quint8] and not post_node_are_quantized:
                         node.output_tensor_infos[0].inf_dtype = node.input_tensor_force_inf_dtype[0]
                         node.insert_fake_quant_after_outputs[0] = True

tests/cpu/test_ao_jit_ipex_quantization.py

@@ -257,8 +257,6 @@ def _lstm_params_list():
                 x = torch.randn(batch_size, seq_len, input_size)
             else:
                 x = torch.randn(seq_len, batch_size, input_size)
-            h = torch.randn(num_layers * num_directions, batch_size, hidden_size)
-            c = torch.randn(num_layers * num_directions, batch_size, hidden_size)
             m = M(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, bidirectional=bidirectional, bias=bias, dropout=dropout, batch_first=batch_first)
             graph = self.checkQuantizeTrace(m, [x], atol=3e-2, rtol=1e-1)
             self.assertGraphContainsExactly(graph, 'ipex::quantized_lstm', 1)
@@ -303,8 +301,10 @@ def forward(self, input, hid=None):
 
         model = M().eval()
         seq = torch.randn(24, 1, 512)
+        h0 = torch.zeros((2, 1, 256), dtype=seq.dtype)
+        hid = (h0, h0)
 
-        graph = self.checkQuantizeTrace(model, [seq], atol=3e-2, rtol=1e-1)
+        graph = self.checkQuantizeTrace(model, [seq, hid], atol=3e-2, rtol=1e-1)
         self.assertGraphContainsExactly(graph, 'ipex::quantized_lstm', 1)        
         self.assertGraphContainsExactly(graph, 'aten::lstm', 0)