Merge branch 'dev' of https://github.com/bayesflow-org/bayesflow into dev

Author: stefanradev93

bayesflow/adapters/transforms/as_set.py

@@ -1,5 +1,5 @@
-import numpy as np
 from keras.saving import register_keras_serializable as serializable
+import numpy as np
 
 from .elementwise_transform import ElementwiseTransform
 

bayesflow/adapters/transforms/sqrt.py

@@ -1,5 +1,5 @@
-import numpy as np
 from keras.saving import register_keras_serializable as serializable
+import numpy as np
 
 from .elementwise_transform import ElementwiseTransform
 

bayesflow/links/ordered.py

@@ -12,6 +12,7 @@ def __init__(self, axis: int, anchor_index: int, **kwargs):
         super().__init__(**keras_kwargs(kwargs))
         self.axis = axis
         self.anchor_index = anchor_index
+        self.group_indices = None
 
         self.config = {"axis": axis, "anchor_index": anchor_index, **kwargs}
 
@@ -22,9 +23,9 @@ def get_config(self):
     def build(self, input_shape):
         super().build(input_shape)
 
-        assert self.anchor_index % input_shape[self.axis] != 0 and self.anchor_index != -1, (
-            "anchor should not be first or last index."
-        )
+        if self.anchor_index % input_shape[self.axis] == 0 or self.anchor_index == -1:
+            raise RuntimeError("Anchor should not be first or last index.")
+
         self.group_indices = dict(
             below=list(range(0, self.anchor_index)),
             above=list(range(self.anchor_index + 1, input_shape[self.axis])),

bayesflow/links/ordered_quantiles.py

@@ -44,16 +44,17 @@ def build(self, input_shape):
         else:
             # choose quantile level closest to median as anchor index
             self.anchor_index = keras.ops.argmin(keras.ops.abs(keras.ops.convert_to_tensor(self.q) - 0.5))
-            msg = (
-                "Length of `q` does not coincide with input shape: "
-                f"len(q)={len(self.q)}, position {self.axis} of shape={input_shape}"
-            )
-            assert num_quantile_levels == len(self.q), msg
 
-        msg = (
-            "The link function `OrderedQuantiles` expects at least 3 quantile levels,"
-            f" but only {num_quantile_levels} were given."
-        )
-        assert self.anchor_index not in (0, -1, num_quantile_levels - 1), msg
+            if len(self.q) != num_quantile_levels:
+                raise RuntimeError(
+                    f"Length of `q` does not coincide with input shape: len(q)={len(self.q)}, "
+                    f"position {self.axis} of shape={input_shape}"
+                )
+
+        if self.anchor_index in [0, -1, num_quantile_levels - 1]:
+            raise RuntimeError(
+                f"The link function `OrderedQuantiles` expects at least 3 quantile levels, "
+                f"but only {num_quantile_levels} were given."
+            )
 
         super().build(input_shape)

bayesflow/links/positive_semi_definite.py

@@ -1,8 +1,10 @@
 import keras
+from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.utils import keras_kwargs
 
 
+@serializable(package="bayesflow.links")
 class PositiveSemiDefinite(keras.Layer):
     """Activation function to link from any square matrix to a positive semidefinite matrix."""
 

bayesflow/networks/consistency_models/consistency_model.py

@@ -178,7 +178,10 @@ def build(self, xz_shape, conditions_shape=None):
         # First, we calculate all unique numbers of discretization steps n
         # in a loop, as self.total_steps might be large
         self.max_n = int(self._schedule_discretization(self.total_steps))
-        assert self.max_n == self.s1 + 1
+
+        if self.max_n != self.s1 + 1:
+            raise ValueError("The maximum number of discretization steps must be equal to s1 + 1.")
+
         unique_n = set()
         for step in range(int(self.total_steps)):
             unique_n.add(int(self._schedule_discretization(step)))

bayesflow/networks/embeddings/fourier_embedding.py

@@ -39,7 +39,10 @@ def __init__(
         """
 
         super().__init__(**kwargs)
-        assert embed_dim % 2 == 0, f"Embedding dimension must be even, but is {embed_dim}."
+
+        if embed_dim % 2 != 0:
+            raise ValueError(f"Embedding dimension must be even, but is {embed_dim}.")
+
         self.w = self.add_weight(initializer=initializer, shape=(embed_dim // 2,), trainable=trainable)
         self.scale = scale
         self.embed_dim = embed_dim

bayesflow/scores/mean_score.py

@@ -1,6 +1,9 @@
+from keras.saving import register_keras_serializable as serializable
+
 from .normed_difference_score import NormedDifferenceScore
 
 
+@serializable(package="bayesflow.scores")
 class MeanScore(NormedDifferenceScore):
     r""":math:`S(\hat \theta, \theta) = | \hat \theta - \theta |^2`
 

bayesflow/scores/median_score.py

@@ -1,6 +1,9 @@
+from keras.saving import register_keras_serializable as serializable
+
 from .normed_difference_score import NormedDifferenceScore
 
 
+@serializable(package="bayesflow.scores")
 class MedianScore(NormedDifferenceScore):
     r""":math:`S(\hat \theta, \theta) = | \hat \theta - \theta |`
 

bayesflow/scores/multivariate_normal_score.py

@@ -1,6 +1,7 @@
 import math
 
 import keras
+from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.links import PositiveSemiDefinite
@@ -9,6 +10,7 @@
 from .parametric_distribution_score import ParametricDistributionScore
 
 
+@serializable(package="bayesflow.scores")
 class MultivariateNormalScore(ParametricDistributionScore):
     r""":math:`S(\hat p_{\mu, \Sigma}, \theta; k) = \log( \mathcal N (\theta; \mu, \Sigma))`
 
@@ -96,9 +98,14 @@ def sample(self, batch_shape: Shape, mean: Tensor, covariance: Tensor) -> Tensor
             A tensor of shape (batch_size, num_samples, D) containing the generated samples.
         """
         batch_size, num_samples = batch_shape
-        dim = mean.shape[-1]
-        assert mean.shape == (batch_size, dim), "mean must have shape (batch_size, D)"
-        assert covariance.shape == (batch_size, dim, dim), "covariance must have shape (batch_size, D, D)"
+        dim = keras.ops.shape(mean)[-1]
+        if keras.ops.shape(mean) != (batch_size, dim):
+            raise ValueError(f"mean must have shape (batch_size, {dim}), but got {keras.ops.shape(mean)}")
+
+        if keras.ops.shape(covariance) != (batch_size, dim, dim):
+            raise ValueError(
+                f"covariance must have shape (batch_size, {dim}, {dim}), but got {keras.ops.shape(covariance)}"
+            )
 
         # Use Cholesky decomposition to generate samples
         cholesky_factor = keras.ops.cholesky(covariance)