fix: gru and lambda layer for tf 2.18&2.19

Author: LongxingTan

.github/workflows/test.yml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, macOS-latest] # add windows-2019 when poetry allows installation with `-f` flag
-        python-version: [3.9, '3.12']
+        python-version: [3.9, '3.11']
         tf-version: [2.13.1, 2.15.1]
 
         exclude:
@@ -92,7 +92,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.12'
+          python-version: '3.11'
 
       - name: Create pip cache directory manually
         run: |

examples/run_prediction_simple.py

@@ -19,7 +19,7 @@
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--seed", type=int, default=315, required=False, help="seed")
-    parser.add_argument("--use_model", type=str, default="rnn", help="model for train")
+    parser.add_argument("--use_model", type=str, default="bert", help="model for train")
     parser.add_argument("--use_data", type=str, default="sine", help="dataset: sine or air passengers")
     parser.add_argument("--train_length", type=int, default=24, help="sequence length for train")
     parser.add_argument("--predict_sequence_length", type=int, default=12, help="sequence length for predict")

tfts/layers/embed_layer.py

@@ -60,9 +60,11 @@ def __init__(self, embed_size: int, positional_type: Optional[str] = "positional
         self.embed_size = embed_size
         self.positional_type = positional_type
 
+    def build(self, input_shape: Tuple[int, ...]):
         # Value embedding layer: the below section is put in init, so it could build while DataEmbedding is call
         # Otherwise, while load the weights, the TokenEmbedding is not built
         self.value_embedding = TokenEmbedding(self.embed_size)
+        self.value_embedding.build(input_shape)
 
         # Positional embedding layer based on specified type
         if self.positional_type == "positional encoding":
@@ -74,6 +76,10 @@ def __init__(self, embed_size: int, positional_type: Optional[str] = "positional
         else:
             self.positional_embedding = None
 
+        if self.positional_embedding:
+            self.positional_embedding.build(input_shape)
+        self.built = True
+
     def call(self, x: tf.Tensor) -> tf.Tensor:
         """
         Forward pass of the layer.

tfts/layers/util_layer.py

@@ -4,8 +4,20 @@
 class ShapeLayer(tf.keras.layers.Layer):
     """Layer to handle shape operations in a Keras-compatible way."""
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
 
     def call(self, x):
         return tf.shape(x)
+
+
+class CreateDecoderFeature(tf.keras.layers.Layer):
+    def __init__(self, predict_sequence_length, **kwargs):
+        super().__init__(**kwargs)
+        self.predict_sequence_length = predict_sequence_length
+
+    def call(self, encoder_feature):
+        batch_size = tf.shape(encoder_feature)[0]
+        time_range = tf.range(self.predict_sequence_length)
+        tiled = tf.tile(tf.reshape(time_range, (1, self.predict_sequence_length, 1)), (batch_size, 1, 1))
+        return tf.cast(tiled, tf.float32)

tfts/models/autoformer.py

@@ -91,11 +91,11 @@ def __init__(self, predict_sequence_length: int = 1, config: Optional[AutoFormer
             hidden_dropout_prob=self.config.hidden_dropout_prob,
         )
 
-        self.project1 = Dense(1, activation=None)
         self.drop1 = Dropout(self.config.hidden_dropout_prob)
         self.dense1 = Dense(512, activation="relu")
         self.drop2 = Dropout(self.config.hidden_dropout_prob)
         self.dense2 = Dense(1024, activation="relu")
+        self.project1 = Dense(1, activation=None)
 
     def __call__(
         self,
@@ -121,7 +121,7 @@ def __call__(
             Otherwise, returns the output tensor.
         """
         x, encoder_feature, decoder_feature = self._prepare_3d_inputs(inputs, ignore_decoder_inputs=False)
-        batch_size, _, n_feature = self.shape_layer(encoder_feature)
+        # batch_size, _, n_feature = self.shape_layer(encoder_feature)
 
         # Encoder
         encoder_output = self.encoder(x)
@@ -198,6 +198,10 @@ def get_config(self):
         base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
+    def compute_output_shape(self, input_shape):
+        batch_size, time_steps, _ = input_shape
+        return (batch_size, time_steps, self.hidden_size)
+
 
 class EncoderLayer(tf.keras.layers.Layer):
     """Encoder Layer for Autoformer architecture."""
@@ -317,6 +321,10 @@ def get_config(self):
         base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
+    def compute_output_shape(self, input_shape):
+        batch_size, time_steps, _ = input_shape
+        return (batch_size, time_steps, self.hidden_size)
+
 
 class DecoderLayer(tf.keras.layers.Layer):
     """Decoder Layer for Autoformer architecture."""

tfts/models/base.py

@@ -11,6 +11,7 @@
 from tensorflow.keras.layers import Concatenate, Lambda
 
 from ..constants import CONFIG_NAME, TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME
+from ..layers.util_layer import CreateDecoderFeature
 
 logger = logging.getLogger(__name__)
 
@@ -95,18 +96,14 @@ def _prepare_3d_inputs(self, inputs, ignore_decoder_inputs=True):
         else:
             encoder_feature = x = inputs
             if not ignore_decoder_inputs:
-                decoder_feature = Lambda(
-                    lambda encoder_feature: tf.cast(
-                        tf.tile(
-                            tf.reshape(tf.range(self.predict_sequence_length), (1, self.predict_sequence_length, 1)),
-                            (tf.shape(encoder_feature)[0], 1, 1),
-                        ),
-                        tf.float32,
-                    ),
-                    output_shape=(self.predict_sequence_length, 1),
-                )(encoder_feature)
+                decoder_feature = CreateDecoderFeature(self.predict_sequence_length)(encoder_feature)
         return x, encoder_feature, decoder_feature
 
+    def _create_decoder_feature(batch_size, predict_sequence_length):
+        time_range = tf.range(predict_sequence_length)
+        tiled = tf.tile(tf.reshape(time_range, (1, predict_sequence_length, 1)), (batch_size, 1, 1))
+        return tf.cast(tiled, tf.float32)
+
     def save_pretrained(
         self,
         save_directory: Union[str, os.PathLike],

tfts/models/rnn.py

@@ -122,6 +122,7 @@ def build(self, input_shape):
                 activation="tanh",
                 return_sequences=True,
                 return_state=return_state,
+                reset_after=False,
                 dropout=self.rnn_dropout if self.rnn_dropout > 0 else 0.0,
             )
 
@@ -213,7 +214,7 @@ def compute_output_shape(self, input_shape):
             elif self.rnn_type == "gru":
                 # GRU: (output, state)
                 return ((batch_size, seq_length, rnn_output_size), (batch_size, rnn_output_size))
-            else:  # LSTM
+            else:
                 # LSTM: (output, state_h, state_c)
                 return ((batch_size, seq_length, rnn_output_size), (batch_size, 2 * rnn_output_size))
 

tfts/models/seq2seq.py

@@ -105,24 +105,35 @@ def __call__(
 class Encoder(tf.keras.layers.Layer):
     def __init__(self, rnn_size, rnn_type="gru", rnn_dropout=0, dense_size=32, return_state=False, **kwargs):
         super().__init__(**kwargs)
+        self.rnn_size = rnn_size
         self.rnn_type = rnn_type.lower()
+        self.rnn_dropout = rnn_dropout
+        self.dense_size = dense_size
         self.return_state = return_state
-        if rnn_type == "gru":
+
+    def build(self, input_shape):
+        if self.rnn_type == "gru":
             self.rnn = GRU(
-                units=rnn_size, activation="tanh", return_state=True, return_sequences=True, dropout=rnn_dropout
+                units=self.rnn_size,
+                activation="tanh",
+                return_state=True,
+                return_sequences=True,
+                dropout=self.rnn_dropout,
+                reset_after=False,
             )
-        elif rnn_type == "lstm":
+        elif self.rnn_type == "lstm":
             self.rnn = LSTM(
-                units=rnn_size,
+                units=self.rnn_size,
                 activation="tanh",
                 return_state=True,
                 return_sequences=True,
-                dropout=rnn_dropout,
+                dropout=self.rnn_dropout,
             )
         else:
-            raise ValueError(f"No supported RNN type: {rnn_type}")
+            raise ValueError(f"No supported RNN type: {self.rnn_type}")
 
-        self.dense = Dense(units=dense_size, activation="tanh")
+        self.dense = Dense(units=self.dense_size, activation="tanh")
+        super(Encoder, self).build(input_shape)
 
     def call(self, inputs):
         """Process input through the encoder RNN and dense layers.
@@ -138,7 +149,8 @@ def call(self, inputs):
             - For LSTM: tuple of (batch_size, dense_size), (batch_size, dense_size)
         """
         if self.rnn_type == "gru":
-            outputs, state = self.rnn(inputs)
+            rnn_outputs = self.rnn(inputs)
+            outputs, state = rnn_outputs
             state = self.dense(state)
         elif self.rnn_type == "lstm":
             outputs, state_h, state_c = self.rnn(inputs)
@@ -151,6 +163,32 @@ def call(self, inputs):
         # outputs = self.dense(outputs)  # => batch_size * input_seq_length * dense_size
         return outputs, state
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "rnn_size": self.rnn_size,
+                "rnn_type": self.rnn_type,
+                "rnn_dropout": self.rnn_dropout,
+                "dense_size": self.dense_size,
+                "return_state": self.return_state,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        batch_size, seq_len, _ = input_shape
+        rnn_output_shape = (batch_size, seq_len, self.rnn_size)
+
+        # State shape depends on RNN type
+        if self.rnn_type == "gru":
+            state_shape = (batch_size, self.dense_size)
+        elif self.rnn_type == "lstm":
+            state_shape = ((batch_size, self.dense_size), (batch_size, self.dense_size))
+        else:
+            raise ValueError(f"No supported rnn type of {self.rnn_type}")
+        return rnn_output_shape, state_shape
+
 
 class DecoderV1(tf.keras.layers.Layer):
     def __init__(
@@ -256,6 +294,29 @@ def call(
         decoder_outputs = tf.concat(decoder_outputs, axis=-1)
         return tf.expand_dims(decoder_outputs, -1)
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "rnn_size": self.rnn_size,
+                "rnn_type": self.rnn_type,
+                "predict_sequence_length": self.predict_sequence_length,
+                "use_attention": self.use_attention,
+                "attention_size": self.attention_size,
+                "num_attention_heads": self.num_attention_heads,
+                "attention_probs_dropout_prob": self.attention_probs_dropout_prob,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        decoder_init_input_shape = input_shape[1]
+        if isinstance(decoder_init_input_shape, (list, tuple)):
+            batch_size = decoder_init_input_shape[0]
+        else:
+            batch_size = None
+        return (batch_size, self.predict_sequence_length, 1)
+
 
 class DecoderV2(tf.keras.layers.Layer):
     def __init__(

tfts/models/transformer.py

@@ -218,16 +218,22 @@ def call(self, inputs: tf.Tensor, mask: Optional[tf.Tensor] = None):
         return x
 
     def get_config(self):
-        config = {
-            "num_hidden_layers": self.num_hidden_layers,
-            "hidden_size": self.hidden_size,
-            "num_attention_heads": self.num_attention_heads,
-            "attention_probs_dropout_prob": self.attention_probs_dropout_prob,
-            "ffn_intermediate_size": self.ffn_intermediate_size,
-            "hidden_dropout_prob": self.hidden_dropout_prob,
-        }
-        base_config = super(Encoder, self).get_config()
-        return dict(list(base_config.items()) + list(config.items()))
+        config = super().get_config()
+        config.update(
+            {
+                "num_hidden_layers": self.num_hidden_layers,
+                "hidden_size": self.hidden_size,
+                "num_attention_heads": self.num_attention_heads,
+                "attention_probs_dropout_prob": self.attention_probs_dropout_prob,
+                "ffn_intermediate_size": self.ffn_intermediate_size,
+                "hidden_dropout_prob": self.hidden_dropout_prob,
+                "layer_norm_eps": self.layer_norm_eps,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
 
 
 class Decoder(tf.keras.layers.Layer):
@@ -332,6 +338,24 @@ def get_causal_attention_mask(self, sequence_length: int) -> tf.Tensor:
         mask = tf.cast(i >= j, dtype="int32")
         return tf.reshape(mask, (1, sequence_length, sequence_length))
 
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_decoder_layers": self.num_decoder_layers,
+                "hidden_size": self.hidden_size,
+                "num_attention_heads": self.num_attention_heads,
+                "attention_probs_dropout_prob": self.attention_probs_dropout_prob,
+                "ffn_intermediate_size": self.ffn_intermediate_size,
+                "hidden_dropout_prob": self.hidden_dropout_prob,
+                "layer_norm_eps": self.layer_norm_eps,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
 
 class DecoderLayer(tf.keras.layers.Layer):
     def __init__(
@@ -399,6 +423,9 @@ def get_config(self):
         base_config = super(DecoderLayer, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
 
 class TransformerBlock(tf.keras.layers.Layer):
     """Basic Transformer block with attention and feed-forward layers."""