update

Author: LongxingTan

tfts/layers/attention_layer.py

@@ -52,6 +52,7 @@ def call(
         training: Optional[bool] = None,
         return_attention_scores: bool = False,
         use_causal_mask: bool = False,
+        **kwargs,
     ):
         """use query and key generating an attention multiplier for value, multi_heads to repeat it
 
@@ -110,12 +111,31 @@ def get_config(self):
         return dict(list(base_config.items()) + list(config.items()))
 
     def compute_output_shape(self, input_shape):
-        if isinstance(input_shape, (list, tuple)) and len(input_shape) == 3:
-            q_shape = input_shape[0]
-        else:
-            raise ValueError("Expected input_shape to be a list or tuple of three elements (q, k, v)")
+        if isinstance(input_shape, tuple) and len(input_shape) == 3:
+            batch_size, seq_len, _ = input_shape
+            return (batch_size, seq_len, self.hidden_size)
+
+        elif isinstance(input_shape, (list, tuple)) and len(input_shape) == 3:
+            q_shape, k_shape, v_shape = input_shape
+
+            # Validate that all shapes are tuples with 3 dimensions
+            if not all(isinstance(shape, tuple) and len(shape) == 3 for shape in [q_shape, k_shape, v_shape]):
+                raise ValueError(
+                    "Each input shape must be a tuple of length 3 (batch_size, seq_len, features). "
+                    f"Got shapes: q={q_shape}, k={k_shape}, v={v_shape}"
+                )
+
+            # Output shape is based on query sequence length
+            batch_size, seq_q_len, _ = q_shape
+            return (batch_size, seq_q_len, self.hidden_size)
 
-        return (q_shape[0], q_shape[1], self.hidden_size)
+        else:
+            raise ValueError(
+                "Expected input_shape to be either:\n"
+                "1. A single tuple (batch_size, seq_len, features) for self-attention, or\n"
+                "2. A list/tuple of 3 shapes [(q_shape), (k_shape), (v_shape)] for cross-attention.\n"
+                f"Got: {input_shape}"
+            )
 
 
 class SelfAttention(tf.keras.layers.Layer):

tfts/models/seq2seq.py

@@ -214,29 +214,35 @@ def __init__(
         self.num_attention_heads = num_attention_heads
         self.attention_probs_dropout_prob = attention_probs_dropout_prob
 
-    def build(self, input_shape):
-        super().build(input_shape)
-        rnn_input_size = input_shape[-1] + 1  # due to in call, concat an initial value
+    def build(self, decoder_features_shape, decoder_init_input_shape, init_state_shape, **kwargs):
+        rnn_input_size = decoder_features_shape[-1] + decoder_init_input_shape[-1]
+
+        if self.use_attention:
+            encoder_output_shape = kwargs.get("encoder_output_shape")
+            if encoder_output_shape is None:
+                raise ValueError("encoder_output_shape must be provided for attention mechanism.")
+            self.attention = Attention(
+                hidden_size=self.attention_size,
+                num_attention_heads=self.num_attention_heads,
+                attention_probs_dropout_prob=self.attention_probs_dropout_prob,
+            )
+            self.attention.build(encoder_output_shape)
+
+            # Add attention output size to RNN input size
+            rnn_input_size += encoder_output_shape[-1]
+
         if self.rnn_type == "gru":
             self.rnn_cell = GRUCell(self.rnn_size)
-            self.rnn_cell.build([None, rnn_input_size])
         elif self.rnn_type == "lstm":
             self.rnn_cell = LSTMCell(units=self.rnn_size)
-            self.rnn_cell.build([None, rnn_input_size])
         else:
-            raise ValueError(f"No supported rnn type of {self.rnn_type}")
+            raise ValueError(f"Unsupported rnn type: {self.rnn_type}")
+
+        self.rnn_cell.build([None, rnn_input_size])
 
         self.dense = Dense(units=1, activation=None)
         self.dense.build([None, self.rnn_size])
-
-        if self.use_attention:
-            self.attention = Attention(
-                hidden_size=self.attention_size,
-                num_attention_heads=self.num_attention_heads,
-                attention_probs_dropout_prob=self.attention_probs_dropout_prob,
-            )
-            self.attention.build(input_shape)
-        self.built = True
+        super().build(decoder_features_shape)
 
     def call(
         self,

tfts/models/wavenet.py

@@ -200,24 +200,31 @@ def __init__(
         self.dilation_rates = dilation_rates
         self.dense_hidden_size = dense_hidden_size
 
-    def build(self, input_shape):
-        super().build(input_shape)
-        batch_size, _, decoder_feature_dim = input_shape
-        decoder_input_size = input_shape[-1] + 1  # due to in call, concat an initial value
+    def build(self, decoder_features_shape, decoder_init_input_shape, encoder_outputs_shape=None, **kwargs):
+        batch_size = decoder_features_shape[0]
+        decoder_input_size = decoder_features_shape[-1] + decoder_init_input_shape[-1]
+
         self.dense1 = Dense(self.filters, activation="tanh")
         self.dense1.build([batch_size, decoder_input_size])
+
         self.dense2 = Dense(2 * self.filters, use_bias=True)
         self.dense2.build([batch_size, self.filters])
+
         self.dense3 = Dense(2 * self.filters, use_bias=False)
         self.dense3.build([batch_size, self.filters])
+
         self.dense4 = Dense(2 * self.filters)
         self.dense4.build([batch_size, self.filters])
-        self.dense5 = Dense(self.dense_hidden_size, activation="relu")
+
         total_skips = self.filters * len(self.dilation_rates)
+        self.dense5 = Dense(self.dense_hidden_size, activation="relu")
         self.dense5.build([batch_size, total_skips])
+
         self.dense6 = Dense(1)
         self.dense6.build([batch_size, self.dense_hidden_size])
 
+        super().build(decoder_features_shape)
+
     def call(
         self,
         decoder_features,