clean up asserts -> raise

also make some code backend agnostic

Author: LarsKue

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 not 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/networks/consistency_models/consistency_model.py

@@ -177,7 +177,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/multivariate_normal_score.py

@@ -98,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)