store summary_dim in summary networks init consistently after output_projector

Author: paul-buerkner

bayesflow/networks/lstnet/lstnet.py

@@ -37,7 +37,6 @@ def __init__(
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.summary_dim = summary_dim
 
         # Convolutional backbone -> can be extended with inception-like structure
         if not isinstance(filters, (list, tuple)):

bayesflow/networks/transformers/fusion_transformer.py

@@ -120,6 +120,7 @@ def __init__(
             raise ValueError("Argument `template_dim` should be in ['lstm', 'gru']")
 
         self.output_projector = keras.layers.Dense(summary_dim)
+        self.summary_dim = summary_dim
 
     def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.

bayesflow/networks/transformers/set_transformer.py

@@ -80,7 +80,6 @@ def __init__(
         """
 
         super().__init__(**kwargs)
-        self.summary_dim = summary_dim
 
         check_lengths_same(embed_dims, num_heads, mlp_depths, mlp_widths)
 
@@ -126,6 +125,7 @@ def __init__(
         )
         self.pooling_by_attention = PoolingByMultiHeadAttention(**(global_attention_settings | pooling_settings))
         self.output_projector = keras.layers.Dense(summary_dim)
+        self.summary_dim = summary_dim
 
     def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Compresses the input sequence into a summary vector of size `summary_dim`.

bayesflow/networks/transformers/time_series_transformer.py

@@ -103,6 +103,7 @@ 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)
+        self.summary_dim = summary_dim
 
         self.time_axis = time_axis