Revert changes in time series / fusion transformers

Author: stefanradev93

bayesflow/networks/coupling_flow/permutations/random.py

@@ -15,14 +15,16 @@ def build(self, xz_shape: Shape, **kwargs) -> None:
 
         self.forward_indices = self.add_weight(
             shape=(xz_shape[-1],),
-            initializer=keras.initializers.Constant(forward_indices),
+            # Best practice: https://github.com/keras-team/keras/pull/20457#discussion_r1832081248
+            initializer=keras.initializers.get(forward_indices),
             trainable=False,
             dtype="int",
         )
 
         self.inverse_indices = self.add_weight(
             shape=(xz_shape[-1],),
-            initializer=keras.initializers.Constant(inverse_indices),
+            # Best practice: https://github.com/keras-team/keras/pull/20457#discussion_r1832081248
+            initializer=keras.initializers.get(inverse_indices),
             trainable=False,
             dtype="int",
         )

bayesflow/networks/transformers/fusion_transformer.py

@@ -2,8 +2,8 @@
 from keras import layers
 
 from bayesflow.types import Tensor
-from bayesflow.utils import check_lengths_same, model_kwargs
-from bayesflow.utils.serialization import deserialize, serializable, serialize
+from bayesflow.utils import check_lengths_same
+from bayesflow.utils.serialization import serializable
 
 from ..summary_network import SummaryNetwork
 
@@ -121,19 +121,6 @@ def __init__(
 
         self.output_projector = keras.layers.Dense(summary_dim)
         self.summary_dim = summary_dim
-        self.embed_dims = embed_dims
-        self.num_heads = num_heads
-        self.mlp_depths = mlp_depths
-        self.mlp_widths = mlp_widths
-        self.dropout = dropout
-        self.mlp_activation = mlp_activation
-        self.kernel_initializer = kernel_initializer
-        self.use_bias = use_bias
-        self.layer_norm = layer_norm
-        self.template_type = template_type
-        self.bidirectional = bidirectional
-        self.template_dim = template_dim
-        self._kwargs = kwargs
 
     def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.
@@ -164,30 +151,3 @@ def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tens
         summary = self.attention_blocks[-1](keras.ops.expand_dims(template, axis=1), rep, training=training, **kwargs)
         summary = self.output_projector(keras.ops.squeeze(summary, axis=1))
         return summary
-
-    @classmethod
-    def from_config(cls, config, custom_objects=None):
-        return cls(**deserialize(config, custom_objects=custom_objects))
-
-    def get_config(self):
-        base_config = super().get_config()
-        base_config = model_kwargs(base_config)
-
-        config = {
-            "summary_dim": self.summary_dim,
-            "embed_dims": self.embed_dims,
-            "num_heads": self.num_heads,
-            "mlp_depths": self.mlp_depths,
-            "mlp_widths": self.mlp_widths,
-            "dropout": self.dropout,
-            "mlp_activation": self.mlp_activation,
-            "kernel_initializer": self.kernel_initializer,
-            "use_bias": self.use_bias,
-            "layer_norm": self.layer_norm,
-            "template_type": self.template_type,
-            "bidirectional": self.bidirectional,
-            "template_dim": self.template_dim,
-            **self._kwargs,
-        }
-
-        return base_config | serialize(config)

bayesflow/networks/transformers/set_transformer.py

@@ -1,8 +1,8 @@
 import keras
 
 from bayesflow.types import Tensor
-from bayesflow.utils import check_lengths_same, model_kwargs
-from bayesflow.utils.serialization import deserialize, serializable, serialize
+from bayesflow.utils import check_lengths_same
+from bayesflow.utils.serialization import serializable
 
 from ..summary_network import SummaryNetwork
 
@@ -129,18 +129,6 @@ def __init__(
         self.output_projector = keras.layers.Dense(summary_dim, name="output_projector")
 
         self.summary_dim = summary_dim
-        self.embed_dims = embed_dims
-        self.num_heads = num_heads
-        self.mlp_depths = mlp_depths
-        self.mlp_widths = mlp_widths
-        self.num_seeds = num_seeds
-        self.dropout = dropout
-        self.mlp_activation = mlp_activation
-        self.kernel_initializer = kernel_initializer
-        self.use_bias = use_bias
-        self.layer_norm = layer_norm
-        self.num_inducing_points = num_inducing_points
-        self.seed_dim = seed_dim
 
     def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.
@@ -165,29 +153,3 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         summary = self.pooling_by_attention(summary, training=training, **kwargs)
         summary = self.output_projector(summary)
         return summary
-
-    @classmethod
-    def from_config(cls, config, custom_objects=None):
-        return cls(**deserialize(config, custom_objects=custom_objects))
-
-    def get_config(self):
-        base_config = super().get_config()
-        base_config = model_kwargs(base_config)
-
-        config = {
-            "summary_dim": self.summary_dim,
-            "embed_dims": self.embed_dims,
-            "num_heads": self.num_heads,
-            "mlp_depths": self.mlp_depths,
-            "mlp_widths": self.mlp_widths,
-            "num_seeds": self.num_seeds,
-            "dropout": self.dropout,
-            "mlp_activation": self.mlp_activation,
-            "kernel_initializer": self.kernel_initializer,
-            "use_bias": self.use_bias,
-            "layer_norm": self.layer_norm,
-            "num_inducing_points": self.num_inducing_points,
-            "seed_dim": self.seed_dim,
-        }
-
-        return base_config | serialize(config)

bayesflow/networks/transformers/time_series_transformer.py

@@ -1,8 +1,8 @@
 import keras
 
 from bayesflow.types import Tensor
-from bayesflow.utils import check_lengths_same, model_kwargs
-from bayesflow.utils.serialization import deserialize, serializable, serialize
+from bayesflow.utils import check_lengths_same
+from bayesflow.utils.serialization import serializable
 
 from ..embeddings import Time2Vec, RecurrentEmbedding
 from ..summary_network import SummaryNetwork
@@ -103,22 +103,9 @@ def __init__(
         # Pooling will be applied as a final step to the abstract representations obtained from set attention
         self.pooling = keras.layers.GlobalAvgPool1D()
         self.output_projector = keras.layers.Dense(summary_dim)
-
-        # store variables for serialization
         self.summary_dim = summary_dim
-        self.embed_dims = embed_dims
-        self.num_heads = num_heads
-        self.mlp_depths = mlp_depths
-        self.mlp_widths = mlp_widths
-        self.dropout = dropout
-        self.mlp_activation = mlp_activation
-        self.kernel_initializer = kernel_initializer
-        self.use_bias = use_bias
-        self.layer_norm = layer_norm
-        self._time_embedding_arg = time_embedding
-        self.time_embed_dim = time_embed_dim
+
         self.time_axis = time_axis
-        self._kwargs = kwargs
 
     def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.
@@ -160,30 +147,3 @@ def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tens
         summary = self.pooling(inp)
         summary = self.output_projector(summary)
         return summary
-
-    @classmethod
-    def from_config(cls, config, custom_objects=None):
-        return cls(**deserialize(config, custom_objects=custom_objects))
-
-    def get_config(self):
-        base_config = super().get_config()
-        base_config = model_kwargs(base_config)
-
-        config = {
-            "summary_dim": self.summary_dim,
-            "embed_dims": self.embed_dims,
-            "num_heads": self.num_heads,
-            "mlp_depths": self.mlp_depths,
-            "mlp_widths": self.mlp_widths,
-            "dropout": self.dropout,
-            "mlp_activation": self.mlp_activation,
-            "kernel_initializer": self.kernel_initializer,
-            "use_bias": self.use_bias,
-            "layer_norm": self.layer_norm,
-            "time_embedding": self._time_embedding_arg,
-            "time_embed_dim": self.time_embed_dim,
-            "time_axis": self.time_axis,
-            **self._kwargs,
-        }
-
-        return base_config | serialize(config)