Remove old usage of aggregate and use weighted_mean everywhere

Author: stefanradev93

bayesflow/experimental/free_form_flow/free_form_flow.py

@@ -12,6 +12,7 @@
  vjp,
  serialize_value_or_type,
  deserialize_value_or_type,
+ weighted_mean,
 )
 
 from bayesflow.networks import InferenceNetwork
@@ -240,6 +241,6 @@ def decode(z):
  reconstruction_loss = ops.sum((x - x_pred) ** 2, axis=-1)
 
  losses = maximum_likelihood_loss + self.beta * reconstruction_loss
- loss = self.aggregate(losses, sample_weight)
+ loss = weighted_mean(losses, sample_weight)
 
  return base_metrics | {"loss": loss}

bayesflow/networks/consistency_models/consistency_model.py

@@ -7,7 +7,13 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, keras_kwargs, serialize_value_or_type, deserialize_value_or_type, weighted_sum
+from bayesflow.utils import (
+ find_network,
+ keras_kwargs,
+ serialize_value_or_type,
+ deserialize_value_or_type,
+ weighted_mean,
+)
 
 
 from ..inference_network import InferenceNetwork
@@ -331,6 +337,6 @@ def compute_metrics(
 
  # Pseudo-huber loss, see [2], Section 3.3
  loss = lam * (ops.sqrt(ops.square(teacher_out - student_out) + self.c_huber2) - self.c_huber)
- loss = weighted_sum(loss, sample_weight)
+ loss = weighted_mean(loss, sample_weight)
 
  return base_metrics | {"loss": loss}

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -7,7 +7,7 @@
  keras_kwargs,
  serialize_value_or_type,
  deserialize_value_or_type,
- weighted_sum,
+ weighted_mean,
 )
 
 from .actnorm import ActNorm
@@ -167,6 +167,6 @@ def compute_metrics(
  base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
 
  z, log_density = self(x, conditions=conditions, inverse=False, density=True)
- loss = weighted_sum(-log_density, sample_weight)
+ loss = weighted_mean(-log_density, sample_weight)
 
  return base_metrics | {"loss": loss}

bayesflow/networks/flow_matching/flow_matching.py

@@ -13,7 +13,7 @@
  optimal_transport,
  serialize_value_or_type,
  deserialize_value_or_type,
- weighted_sum,
+ weighted_mean,
 )
 from ..inference_network import InferenceNetwork
 
@@ -260,6 +260,6 @@ def compute_metrics(
  predicted_velocity = self.velocity(x, time=t, conditions=conditions, training=stage == "training")
 
  loss = self.loss_fn(target_velocity, predicted_velocity)
- loss = weighted_sum(loss, sample_weight)
+ loss = weighted_mean(loss, sample_weight)
 
  return base_metrics | {"loss": loss}

bayesflow/scores/normed_difference_score.py

@@ -2,6 +2,7 @@
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
+from bayesflow.utils import weighted_mean
 
 from .scoring_rule import ScoringRule
 
@@ -55,7 +56,7 @@ def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor =
  """
  estimates = estimates["value"]
  scores = keras.ops.absolute(estimates - targets) ** self.k
- score = self.aggregate(scores, weights)
+ score = weighted_mean(scores, weights)
  return score
 
  def get_config(self):

bayesflow/scores/parametric_distribution_score.py

@@ -1,6 +1,7 @@
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
+from bayesflow.utils import weighted_mean
 
 from .scoring_rule import ScoringRule
 
@@ -29,5 +30,5 @@ def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor =
  :math:`S(\hat p_\phi, \theta; k) = -\log(\hat p_\phi(\theta))`
  """
  scores = -self.log_prob(x=targets, **estimates)
- score = self.aggregate(scores, weights)
+ score = weighted_mean(scores, weights)
  return score

bayesflow/scores/quantile_score.py

@@ -4,7 +4,7 @@
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
-from bayesflow.utils import logging
+from bayesflow.utils import logging, weighted_mean
 from bayesflow.links import OrderedQuantiles
 
 from .scoring_rule import ScoringRule
@@ -39,7 +39,7 @@ def get_config(self):
  base_config = super().get_config()
  return base_config | self.config
 
- def get_head_shapes_from_target_shape(self, target_shape: Shape):
+ def get_head_shapes_from_target_shape(self, target_shape: Shape) -> dict[str, tuple]:
  # keras.saving.load_model sometimes passes target_shape as a list, so we force a conversion
  target_shape = tuple(target_shape)
  return dict(value=(len(self.q),) + target_shape[1:])
@@ -49,5 +49,5 @@ def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor =
  pointwise_differance = estimates - targets[:, None, :]
  scores = pointwise_differance * (keras.ops.cast(pointwise_differance > 0, float) - self._q[None, :, None])
  scores = keras.ops.mean(scores, axis=1)
- score = self.aggregate(scores, weights)
+ score = weighted_mean(scores, weights)
  return score

bayesflow/scores/scoring_rule.py

@@ -204,30 +204,3 @@ def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor)
  <tf.Tensor: shape=(), dtype=float32, numpy=1.013983130455017>
  """
  raise NotImplementedError
-
- def aggregate(self, scores: Tensor, weights: Tensor = None) -> Tensor:
- """
- Computes the mean of **scores**, optionally applying **weights**.
-
- This function computes the mean value of the given scores. When weights are provided,
- it first multiplies the scores by the weights and then computes the mean of the result.
- If no weights are provided, it computes the mean of the scores.
-
- Parameters
- ----------
- scores : Tensor
- A tensor containing the scores to be aggregated.
- weights : Tensor, optional (default - None)
- A tensor of weights corresponding to each score. Must be the same shape as `scores`.
- If not provided, the function returns the mean of `scores`.
-
- Returns
- -------
- Tensor
- The aggregated score computed as a weighted mean if **weights** is provided,
- or as the simple mean of **scores** otherwise.
- """
-
- if weights is not None:
- return keras.ops.mean(scores * weights)
- return keras.ops.mean(scores)

bayesflow/utils/__init__.py

@@ -71,7 +71,7 @@
  tree_concatenate,
  tree_stack,
  fill_triangular_matrix,
- weighted_sum,
+ weighted_mean,
 )
 from .classification import calibration_curve, confusion_matrix
 from .validators import check_lengths_same

bayesflow/utils/tensor_utils.py

@@ -140,7 +140,7 @@ def pad(x: Tensor, value: float | Tensor, n: int, axis: int, side: str = "both")
  raise TypeError(f"Invalid side type {type(side)!r}. Must be str.")
 
 
-def weighted_sum(elements: Tensor, weights: Tensor = None) -> Tensor:
+def weighted_mean(elements: Tensor, weights: Tensor = None) -> Tensor:
  """
  Compute the (optionally) weighted mean of the input tensor.