Hotfix: numercial stability of non-log-stabilized sinkhorn plan

[LarsKue]

This is an important hotfix, so sending it directly to main. Thanks @daniel-habermann for the report!

I also did some benchmarking and it seems the current implementation is optimal with respect to performance.

Edit: I also fixed an issue with the convergence check in log_sinkhorn_plan which takes its runtime down from ~4s to ~0.02s.

[Copilot]

Pull Request Overview

This hotfix ensures numerical stability for the non-log-stabilized Sinkhorn implementation and aligns tests to cover both the new and existing normalization methods.

• Updated the Sinkhorn plan initialization and normalization logic for better stability.
• Changed max_steps default to None to run until convergence.
• Extended existing tests to parameterize over "log_sinkhorn" and "sinkhorn" methods.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
tests/test_utils/test_optimal_transport.py	Parameterized two tests over both Sinkhorn variants and removed a skip decorator
bayesflow/utils/optimal_transport/sinkhorn.py	Added numerical-stability steps, switched normalization from softmax to sum-based, and updated max_steps default

Comments suppressed due to low confidence (1)

bayesflow/utils/optimal_transport/sinkhorn.py:45

• The default value of None conflicts with the int annotation. Consider changing the signature to max_steps: Optional[int] = None and importing Optional from typing for clarity.

    max_steps: int = None,

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Files with missing lines	Coverage Δ
bayesflow/diagnostics/metrics/__init__.py	100.00% <100.00%> (ø)
...sflow/diagnostics/metrics/calibration_log_gamma.py	100.00% <100.00%> (ø)
bayesflow/distributions/diagonal_normal.py	95.34% <100.00%> (-0.11%)	⬇️
bayesflow/distributions/diagonal_student_t.py	95.91% <100.00%> (-0.09%)	⬇️
bayesflow/distributions/mixture.py	98.03% <ø> (ø)
...esflow/networks/standardization/standardization.py	95.45% <100.00%> (ø)
bayesflow/scores/multivariate_normal_score.py	97.72% <100.00%> (ø)
bayesflow/utils/optimal_transport/log_sinkhorn.py	100.00% <100.00%> (ø)
bayesflow/utils/optimal_transport/sinkhorn.py	100.00% <100.00%> (ø)

... and 2 files with indirect coverage changes

[daniel-habermann]

I have added some additional commits. Most notable changes are:

• sinkhorn_plan and log_sinkhorn_plan now return proper transport plans, that is, the row and column sums match the marginals.
• the assignments in sinkhorn and log_sinkhorn are now calculated as assignments = keras.ops.categorical(log_plan), because keras.ops.categorical expects log-probs instead of probabilities as inputs (which they confusingly call logits).
• added some additional unit tests to ensure that the transport plans are correct.

@LarsKue could you please double check that the comment in this function definition is what you intended? (the such that.. part)

def sinkhorn(x1: Tensor, x2: Tensor, seed: int = None, **kwargs) -> (Tensor, Tensor):
    """
    Matches elements from x2 onto x1 using the Sinkhorn-Knopp algorithm.

    Sinkhorn-Knopp is an iterative algorithm that repeatedly normalizes the cost matrix into a
    transport plan, containing assignment probabilities.
    The permutation is then sampled randomly according to the transport plan.

    :param x1: Tensor of shape (n, ...)
        Samples from the first distribution.

    :param x2: Tensor of shape (m, ...)
        Samples from the second distribution.

    :param kwargs:
        Additional keyword arguments that are passed to :py:func:`sinkhorn_plan`.

    :param seed: Random seed to use for sampling indices.
        Default: None, which means the seed will be auto-determined for non-compiled contexts.

    :return: Tensor of shape (n,)
        Assignment indices for x2.

    """
    plan = sinkhorn_plan(x1, x2, **kwargs)

    # we sample from log(plan) to receive assignments of length n, corresponding to indices of x2
    # such that x2[assignments] matches x1
    assignments = keras.random.categorical(keras.ops.log(plan), num_samples=1, seed=seed)
    assignments = keras.ops.squeeze(assignments, axis=1)

    return assignments

From my point of view this is ready to be merged now.