Merge branch 'dev' into c2st

Author: stefanradev93

.github/workflows/multiversion-docs.yaml

@@ -43,7 +43,7 @@ jobs:
         run: |
           cd ./repo/docsrc
           make clean
-          make docs-sequential
+          make production-docs-sequential
 
       - name: Checkout gh-pages-dev
         uses: actions/checkout@v3

.github/workflows/test-docs.yaml

@@ -31,7 +31,7 @@ jobs:
         run: |
           cd ./docsrc
           make clean
-          make local
+          make local-docs
       - name: Clean up
         run: |
           cd ./docsrc

CONTRIBUTING.md

@@ -171,17 +171,17 @@ You can re-build the for your local state with:
 
 ```bash
 cd docsrc
-make clean && make local
+make clean && make local-docs
 # in case of issues, try `make clean-all`
 ```
 
 Note that files ignored by git (i.e., listed in `.gitignore`) are not included in the documentation.
 
-We also provide a multi-version documentation. To generate it, run
+We also provide a multi-version documentation, which renders the branches `main` and `stable-legacy`. To generate it, run
 
 ```bash
 cd docsrc
-make clean && make docs
+make clean && make production-docs
 ```
 
 This will create and cache virtual environments for the build at `docsrc/.docs_venvs`.

README.md

@@ -63,12 +63,16 @@ More tutorials are always welcome! Please consider making a pull request if you
 
 ## Install
 
-BayesFlow is available to install via pip:
+BayesFlow v2 is not yet installable via PyPI, but you can use the following command to install the latest version of the `main` branch:
 
 ```bash
-pip install bayesflow
+pip install git+https://github.com/bayesflow-org/bayesflow.git
 ```
 
+If you encounter problems with this or require more control, please refer to the instructions to install from source below.
+
+Note: `pip install bayesflow` will install the v1 version of BayesFlow.
+
 ### Backend
 
 To use BayesFlow, you will also need to install one of the following machine learning backends.

bayesflow/networks/consistency_models/consistency_model.py

@@ -7,7 +7,7 @@
 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
+from bayesflow.utils import find_network, keras_kwargs, serialize_value_or_type, deserialize_value_or_type, weighted_sum
 
 
 from ..inference_network import InferenceNetwork
@@ -285,7 +285,9 @@ def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None,
         out = skip * x + out * f
         return out
 
-    def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "training") -> dict[str, Tensor]:
+    def compute_metrics(
+        self, x: Tensor, conditions: Tensor = None, sample_weight: Tensor = None, stage: str = "training"
+    ) -> dict[str, Tensor]:
         base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
 
         # The discretization schedule requires the number of passed training steps.
@@ -328,6 +330,7 @@ def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "tr
         lam = 1 / (t2 - t1)
 
         # Pseudo-huber loss, see [2], Section 3.3
-        loss = ops.mean(lam * (ops.sqrt(ops.square(teacher_out - student_out) + self.c_huber2) - self.c_huber))
+        loss = lam * (ops.sqrt(ops.square(teacher_out - student_out) + self.c_huber2) - self.c_huber)
+        loss = weighted_sum(loss, sample_weight)
 
         return base_metrics | {"loss": loss}

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -2,7 +2,13 @@
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_permutation, keras_kwargs, serialize_value_or_type, deserialize_value_or_type
+from bayesflow.utils import (
+    find_permutation,
+    keras_kwargs,
+    serialize_value_or_type,
+    deserialize_value_or_type,
+    weighted_sum,
+)
 
 from .actnorm import ActNorm
 from .couplings import DualCoupling
@@ -158,11 +164,9 @@ def _inverse(
     def compute_metrics(
         self, x: Tensor, conditions: Tensor = None, sample_weight: Tensor = None, stage: str = "training"
     ) -> dict[str, Tensor]:
-        if sample_weight is not None:
-            print(sample_weight)
-        base_metrics = super().compute_metrics(x, conditions=conditions, sample_weight=sample_weight, stage=stage)
+        base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
 
         z, log_density = self(x, conditions=conditions, inverse=False, density=True)
-        loss = self.aggregate(-log_density, sample_weight)
+        loss = weighted_sum(-log_density, sample_weight)
 
         return base_metrics | {"loss": loss}

bayesflow/networks/flow_matching/flow_matching.py

@@ -13,6 +13,7 @@
     optimal_transport,
     serialize_value_or_type,
     deserialize_value_or_type,
+    weighted_sum,
 )
 from ..inference_network import InferenceNetwork
 
@@ -254,11 +255,11 @@ def compute_metrics(
             x = t * x1 + (1 - t) * x0
             target_velocity = x1 - x0
 
-        base_metrics = super().compute_metrics(x1, conditions, sample_weight, stage)
+        base_metrics = super().compute_metrics(x1, conditions=conditions, stage=stage)
 
         predicted_velocity = self.velocity(x, time=t, conditions=conditions, training=stage == "training")
 
         loss = self.loss_fn(target_velocity, predicted_velocity)
-        loss = self.aggregate(loss, sample_weight)
+        loss = weighted_sum(loss, sample_weight)
 
         return base_metrics | {"loss": loss}

bayesflow/networks/inference_network.py

@@ -48,9 +48,7 @@ def log_prob(self, samples: Tensor, conditions: Tensor = None, **kwargs) -> Tens
         _, log_density = self(samples, conditions=conditions, inverse=False, density=True, **kwargs)
         return log_density
 
-    def compute_metrics(
-        self, x: Tensor, conditions: Tensor = None, sample_weight: Tensor = None, stage: str = "training"
-    ) -> dict[str, Tensor]:
+    def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "training") -> dict[str, Tensor]:
         if not self.built:
             xz_shape = keras.ops.shape(x)
             conditions_shape = None if conditions is None else keras.ops.shape(conditions)
@@ -66,10 +64,3 @@ def compute_metrics(
                 metrics[metric.name] = metric(samples, x)
 
         return metrics
-
-    def aggregate(self, losses: Tensor, weights: Tensor = None):
-        if weights is not None:
-            weighted = losses * weights
-        else:
-            weighted = losses
-        return keras.ops.mean(weighted)

bayesflow/utils/__init__.py

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

bayesflow/utils/tensor_utils.py

@@ -140,6 +140,26 @@ 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:
+    """
+    Compute the (optionally) weighted mean of the input tensor.
+
+    Parameters
+    ----------
+    elements : Tensor
+        A tensor containing the elements to average.
+    weights : Tensor, optional
+        A tensor of the same shape as `elements` representing weights.
+        If None, the mean is computed without weights.
+
+    Returns
+    -------
+    Tensor
+        A scalar tensor representing the (weighted) mean.
+    """
+    return keras.ops.mean(elements * weights if weights is not None else elements)
+
+
 def searchsorted(sorted_sequence: Tensor, values: Tensor, side: str = "left") -> Tensor:
     """
     Find indices where elements should be inserted to maintain order.