Merge branch 'dev' into keras-3.9

Author: LarsKue

.gitignore

@@ -10,6 +10,7 @@ docsrc/.docs_venvs
 docsrc/source/api
 docsrc/source/_examples
 docsrc/source/contributing.md
+examples/checkpoints/
 build
 docs/
 

.pre-commit-config.yaml

@@ -14,7 +14,7 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
   # Ruff version.
-  rev: v0.9.6
+  rev: v0.9.10
   hooks:
     # Run the linter.
     - id: ruff

bayesflow/diagnostics/plots/z_score_contraction.py

@@ -121,7 +121,7 @@ def z_score_contraction(
         ax.scatter(contraction[:, i], z_score[:, i], color=color, alpha=0.5)
         ax.set_xlim([-0.05, 1.05])
 
-    prettify_subplots(plot_data["axes"], tick_fontsize)
+    prettify_subplots(plot_data["axes"], num_subplots=plot_data["num_variables"], tick_fontsize=tick_fontsize)
 
     # Add labels, titles, and set font sizes
     add_titles_and_labels(

bayesflow/networks/transformers/fusion_transformer.py

@@ -4,6 +4,7 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
+from bayesflow.utils.decorators import sanitize_input_shape
 
 from ..summary_network import SummaryNetwork
 
@@ -151,3 +152,8 @@ 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
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.call(keras.ops.zeros(input_shape))

bayesflow/networks/transformers/set_transformer.py

@@ -3,6 +3,7 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
+from bayesflow.utils.decorators import sanitize_input_shape
 
 from ..summary_network import SummaryNetwork
 
@@ -150,3 +151,8 @@ 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
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.call(keras.ops.zeros(input_shape))

bayesflow/networks/transformers/time_series_transformer.py

@@ -3,6 +3,7 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import check_lengths_same
+from bayesflow.utils.decorators import sanitize_input_shape
 
 from ..embeddings import Time2Vec, RecurrentEmbedding
 from ..summary_network import SummaryNetwork
@@ -147,3 +148,8 @@ def call(self, input_sequence: Tensor, training: bool = False, **kwargs) -> Tens
         summary = self.pooling(inp)
         summary = self.output_projector(summary)
         return summary
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        super().build(input_shape)
+        self.call(keras.ops.zeros(input_shape))

bayesflow/utils/plot_utils.py

@@ -203,8 +203,8 @@ def add_y_labels(axes: np.ndarray, num_row: int = None, ylabel: Sequence[str] |
 
 
 def add_titles(axes: np.ndarray, title: Sequence[str] | str = None, title_fontsize: int = None):
-    for i, ax in enumerate(axes.flat):
-        ax.set_title(title[i], fontsize=title_fontsize)
+    for t, ax in zip(title, axes.flat):
+        ax.set_title(t, fontsize=title_fontsize)
 
 
 def add_titles_and_labels(