Merge branch 'dev' into feat-test-non-invertible

Author: vpratz

.github/workflows/publish.yaml

@@ -1,18 +1,53 @@
+name: Publish Python 🐍 distribution 📦 to PyPI
 
-name: Publish to PyPI.org
 on:
   release:
     types: [published]
+
 jobs:
-  pypi:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        persist-credentials: false
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+    - name: Install pypa/build
+      run: >-
+        python3 -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    needs:
+    - build
     runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/bayesflow  # Replace <package-name> with your PyPI project name
+    permissions:
+      id-token: write  # IMPORTANT: mandatory for trusted publishing
+
     steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-      - run: python3 -m pip install -U build && python3 -m build
-      - name: Publish package
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          password: ${{ secrets.PYPI_API_TOKEN }}
+    - name: Download all the dists
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1

README.md

@@ -64,15 +64,19 @@ More tutorials are always welcome! Please consider making a pull request if you
 
 ## Install
 
-BayesFlow v2 is not yet installable via PyPI, but you can use the following command to install the latest version of the `main` branch:
+You can install the latest stable version from PyPI using:
 
 ```bash
-pip install git+https://github.com/bayesflow-org/bayesflow.git
+pip install bayesflow
 ```
 
-If you encounter problems with this or require more control, please refer to the instructions to install from source below.
+If you want the latest features, you can install from source:
 
-Note: `pip install bayesflow` will install the v1 version of BayesFlow.
+```bash
+pip install git+https://github.com/bayesflow-org/bayesflow.git@dev
+```
+
+If you encounter problems with this or require more control, please refer to the instructions to install from source below.
 
 ### Backend
 

bayesflow/diagnostics/plots/loss.py

@@ -19,6 +19,8 @@ def loss(
     figsize: Sequence[float] = None,
     train_color: str = "#132a70",
     val_color: str = "black",
+    val_marker: str = "o",
+    val_marker_size: float = 5,
     lw_train: float = 2.0,
     lw_val: float = 2.0,
     grid_alpha: float = 0.2,
@@ -49,10 +51,14 @@ def loss(
         The color for the train loss trajectory
     val_color          : str, optional, default: None
         The color for the optional validation loss trajectory
+    val_marker: str
+        Marker style for the validation loss curve. Default is "o".
+    val_marker_size: float
+        Marker size for the validation loss curve. Default is 5.
     lw_train           : int, optional, default: 2
-        The linewidth for the training loss curve
+        The line width for the training loss curve
     lw_val             : int, optional, default: 2
-        The linewidth for the validation loss curve
+        The line width for the validation loss curve
     grid_alpha          : float, optional, default: 0.2
         The transparency of the background grid
     legend_fontsize    : int, optional, default: 14
@@ -130,6 +136,9 @@ def loss(
                 color=val_color,
                 lw=lw_val,
                 alpha=alpha_unsmoothed,
+                linestyle="--",
+                marker=val_marker,
+                markersize=val_marker_size,
                 label="Validation",
             )
 
@@ -140,6 +149,7 @@ def loss(
                     val_step_index,
                     smoothed_val_loss,
                     color=val_color,
+                    linestyle="--",
                     lw=lw_val,
                     alpha=0.8,
                     label="Validation (Moving Average)",

bayesflow/networks/flow_matching/flow_matching.py

@@ -82,7 +82,7 @@ def __init__(
             The base probability distribution from which samples are drawn, such as "normal".
             Default is "normal".
         use_optimal_transport : bool, optional
-            Whether to apply optimal transport for improved training stability. Default is False.
+            Whether to apply optimal transport for improved training stability. Default is True.
         loss_fn : str, optional
             The loss function used for training, such as "mse". Default is "mse".
         integrate_kwargs : dict[str, any], optional
@@ -256,9 +256,20 @@ def compute_metrics(
             x0 = self.base_distribution.sample(keras.ops.shape(x1)[:-1])
 
             if self.use_optimal_transport:
-                x1, x0, conditions = optimal_transport(
-                    x1, x0, conditions, seed=self.seed_generator, **self.optimal_transport_kwargs
+                # we must choose between resampling x0 or x1
+                # since the data is possibly noisy and may contain outliers, it is better
+                # to possibly drop some samples from x1 than from x0
+                # in the marginal over multiple batches, this is not a problem
+                x0, x1, assignments = optimal_transport(
+                    x0,
+                    x1,
+                    seed=self.seed_generator,
+                    **self.optimal_transport_kwargs,
+                    return_assignments=True,
                 )
+                if conditions is not None:
+                    # conditions must be resampled along with x1
+                    conditions = keras.ops.take(conditions, assignments, axis=0)
 
             t = keras.random.uniform((keras.ops.shape(x0)[0],), seed=self.seed_generator)
             t = expand_right_as(t, x0)

bayesflow/utils/optimal_transport/__init__.py

@@ -1,2 +1 @@
 from .optimal_transport import optimal_transport
-from .sinkhorn import sinkhorn, sinkhorn_indices, sinkhorn_plan

bayesflow/utils/optimal_transport/euclidean.py

@@ -0,0 +1,11 @@
+import keras
+
+
+def euclidean(x1, x2):
+    # TODO: rename and move this function
+    result = x1[:, None] - x2[None, :]
+    shape = list(keras.ops.shape(result))
+    shape[2:] = [-1]
+    result = keras.ops.reshape(result, shape)
+    result = keras.ops.norm(result, ord=2, axis=-1)
+    return result

bayesflow/utils/optimal_transport/hungarian.py

@@ -0,0 +1,80 @@
+import keras
+
+from .. import logging
+from ..tensor_utils import is_symbolic_tensor
+
+from .euclidean import euclidean
+
+
+def log_sinkhorn(x1, x2, seed: int = None, **kwargs):
+    """
+    Log-stabilized version of :py:func:`~bayesflow.utils.optimal_transport.sinkhorn.sinkhorn`.
+    Significantly slower than the unstabilized version, so use only when you need numerical stability.
+    """
+    log_plan = log_sinkhorn_plan(x1, x2, **kwargs)
+    assignments = keras.random.categorical(keras.ops.exp(log_plan), num_samples=1, seed=seed)
+    assignments = keras.ops.squeeze(assignments, axis=1)
+
+    return assignments
+
+
+def log_sinkhorn_plan(x1, x2, regularization: float = 1.0, rtol=1e-5, atol=1e-8, max_steps=None):
+    """
+    Log-stabilized version of :py:func:`~bayesflow.utils.optimal_transport.sinkhorn.sinkhorn_plan`.
+    Significantly slower than the unstabilized version, so use only when you need numerical stability.
+    """
+    cost = euclidean(x1, x2)
+
+    log_plan = cost / -(regularization * keras.ops.mean(cost) + 1e-16)
+
+    if is_symbolic_tensor(log_plan):
+        return log_plan
+
+    def contains_nans(plan):
+        return keras.ops.any(keras.ops.isnan(plan))
+
+    def is_converged(plan):
+        # for convergence, the plan should be doubly stochastic
+        conv0 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=0), 0.0, rtol=rtol, atol=atol))
+        conv1 = keras.ops.all(keras.ops.isclose(keras.ops.logsumexp(plan, axis=1), 0.0, rtol=rtol, atol=atol))
+        return conv0 & conv1
+
+    def cond(_, plan):
+        # break the while loop if the plan contains nans or is converged
+        return ~(contains_nans(plan) | is_converged(plan))
+
+    def body(steps, plan):
+        # Sinkhorn-Knopp: repeatedly normalize the transport plan along each dimension
+        plan = keras.ops.log_softmax(plan, axis=0)
+        plan = keras.ops.log_softmax(plan, axis=1)
+
+        return steps + 1, plan
+
+    steps = 0
+    steps, log_plan = keras.ops.while_loop(cond, body, (steps, log_plan), maximum_iterations=max_steps)
+
+    def do_nothing():
+        pass
+
+    def log_steps():
+        msg = "Log-Sinkhorn-Knopp converged after {:d} steps."
+
+        logging.info(msg, steps)
+
+    def warn_convergence():
+        marginals = keras.ops.logsumexp(log_plan, axis=0)
+        deviations = keras.ops.abs(marginals)
+        badness = 100.0 * keras.ops.exp(keras.ops.max(deviations))
+
+        msg = "Log-Sinkhorn-Knopp did not converge after {:d} steps (badness: {:.1f}%)."
+
+        logging.warning(msg, max_steps, badness)
+
+    def warn_nans():
+        msg = "Log-Sinkhorn-Knopp produced NaNs."
+        logging.warning(msg)
+
+    keras.ops.cond(contains_nans(log_plan), warn_nans, do_nothing)
+    keras.ops.cond(is_converged(log_plan), log_steps, warn_convergence)
+
+    return log_plan

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -1,13 +1,17 @@
-from bayesflow.types import Tensor
+import keras
 
-from .hungarian import hungarian
-from .random import random
+from .log_sinkhorn import log_sinkhorn
 from .sinkhorn import sinkhorn
 
+methods = {
+    "sinkhorn": sinkhorn,
+    "sinkhorn_knopp": sinkhorn,
+    "log_sinkhorn": log_sinkhorn,
+    "log_sinkhorn_knopp": log_sinkhorn,
+}
 
-def optimal_transport(
-    x1: Tensor, x2: Tensor, *aux: Tensor, method: str = "sinkhorn_knopp", **kwargs
-) -> (Tensor, Tensor):
+
+def optimal_transport(x1, x2, method="log_sinkhorn", return_assignments=False, **kwargs):
     """Matches elements from x2 onto x1, such that the transport cost between them is minimized, according to the method
     and cost matrix used.
 
@@ -23,28 +27,21 @@ def optimal_transport(
     :param x2: Tensor of shape (m, ...)
         Samples from the second distribution.
 
-    :param aux: Tensors of shape (n, ...)
-        Auxiliary tensors to be permuted along with x1.
-        Note that x2 is never permuted for all currently available methods.
-
     :param method: Method used to compute the transport cost.
-        Default: 'sinkhorn_knopp'
+        Default: 'log_sinkhorn'
 
-    :param kwargs: Additional keyword arguments passed to the optimization method.
+    :param return_assignments: Whether to return the assignment indices.
+        Default: False
+
+    :param kwargs: Additional keyword arguments that are passed to the optimization method.
 
     :return: Tensors of shapes (n, ...) and (m, ...)
         x1 and x2 in optimal transport permutation order.
     """
-    methods = {
-        "hungarian": hungarian,
-        "sinkhorn": sinkhorn,
-        "sinkhorn_knopp": sinkhorn,
-        "random": random,
-    }
-
-    method = method.lower()
+    assignments = methods[method.lower()](x1, x2, **kwargs)
+    x2 = keras.ops.take(x2, assignments, axis=0)
 
-    if method not in methods:
-        raise ValueError(f"Unsupported method name: '{method}'.")
+    if return_assignments:
+        return x1, x2, assignments
 
-    return methods[method](x1, x2, *aux, **kwargs)
+    return x1, x2