Revert "Simplify multi headed attention (#1153)" (#1155)

This reverts commit 3d9a418.

onmt/modules/multi_headed_attn.py

@@ -120,16 +120,14 @@ def unshape(x):
             return x.transpose(1, 2).contiguous() \
                     .view(batch_size, -1, head_count * dim_per_head)
 
-        query, key, value = self.linear_query(query),\
-            self.linear_keys(key),\
-            self.linear_values(value)
-        key = shape(key)
-        value = shape(value)
-        query = shape(query)
-
         # 1) Project key, value, and query.
         if layer_cache is not None:
             if type == "self":
+                query, key, value = self.linear_query(query),\
+                                    self.linear_keys(query),\
+                                    self.linear_values(query)
+                key = shape(key)
+                value = shape(value)
                 device = key.device
                 if layer_cache["self_keys"] is not None:
                     key = torch.cat(
@@ -142,11 +140,25 @@ def unshape(x):
                 layer_cache["self_keys"] = key
                 layer_cache["self_values"] = value
             elif type == "context":
-                if layer_cache["memory_keys"] is not None:
+                query = self.linear_query(query)
+                if layer_cache["memory_keys"] is None:
+                    key, value = self.linear_keys(key),\
+                                 self.linear_values(value)
+                    key = shape(key)
+                    value = shape(value)
+                else:
                     key, value = layer_cache["memory_keys"],\
                                layer_cache["memory_values"]
                 layer_cache["memory_keys"] = key
                 layer_cache["memory_values"] = value
+        else:
+            key = self.linear_keys(key)
+            value = self.linear_values(value)
+            query = self.linear_query(query)
+            key = shape(key)
+            value = shape(value)
+
+        query = shape(query)
 
         key_len = key.size(2)
         query_len = query.size(2)