LSTM bf16: support bias and c_states to be bf16 (#154)

* use bf16 bias, cx and cy of bf16 LSTM

* clang format

Author: chunyuan-w

tests/cpu/test_autocast.py

@@ -272,14 +272,12 @@ def _test_lstm(self, training, bf16, prec = 1e-5):
                     self.assertEqual(h_ipex.dtype, torch.float)
                     self.assertEqual(c_ipex.dtype, torch.float)
 
-                    # with mkldnn LSTM, y, hy[0] is bf16 and hy[1] is fp32
                     self.assertEqual(y_ipex.dtype, torch.bfloat16)
                     self.assertEqual(hy_ipex[0].dtype, torch.bfloat16)
-                    self.assertEqual(hy_ipex[1].dtype, torch.float)
+                    self.assertEqual(hy_ipex[1].dtype, torch.bfloat16)
                 self.assertEqual(y, y_ipex, prec=prec)
                 self.assertEqual(hy[0], hy_ipex[0], prec=prec)
-
-                self.assertEqual(hy[1], self._cast_dtype(hy_ipex[1], bf16), prec=prec)
+                self.assertEqual(hy[1], hy_ipex[1], prec=prec)
 
     def _test_lstm_pack_padded_sequence(self):
         embedding_dim = 1024

tests/cpu/test_rnnt_custom_kernel.py

@@ -120,8 +120,8 @@ def test_rnnt_update_batch(self):
 
             x_org = torch.randn([self.max_len, batch_size, 2], dtype=dtype)
             x = copy.deepcopy(x_org)
-            hidden = [torch.zeros([2, batch_size, 320], dtype=dtype), torch.zeros([2, batch_size, 320], dtype=torch.float)]
-            hidden_prime = [torch.randn([2, batch_size, 320], dtype=dtype), torch.randn([2, batch_size, 320], dtype=torch.float)]
+            hidden = [torch.zeros([2, batch_size, 320], dtype=dtype), torch.zeros([2, batch_size, 320], dtype=dtype)]
+            hidden_prime = [torch.randn([2, batch_size, 320], dtype=dtype), torch.randn([2, batch_size, 320], dtype=dtype)]
 
             blank_vec_org, blankness_org, label_col_org, time_idxs_org, symbols_added_org, not_blank_org, label_tensor_org, hidden_org, f_org = self._test_org(hidden, hidden_prime, x_org.transpose(0, 1), batch_size, max_symbol, blank_id, loop_cnt)
             blank_vec_out, blankness_out, label_col, time_idxs, symbols_added, not_blank, label_tensor, hidden, f = self._test_rnnt_update_batch_kernel(hidden, hidden_prime, x.transpose(0,1), batch_size, max_symbol, blank_id, loop_cnt)