[WIP] Generalize convolutional Wasserstein barycenters to N-D grids

[tvercaut]

Summary

Pure refactor of ot/bregman/_convolutional.py: no new solver, no new maths.

• Commit 1 fixes a latent API-conformance bug in Backend.logsumexp (Cupy/Tensorflow didn't accept keepdims, unlike NumPy/Jax/Torch) — a prerequisite for writing backend-agnostic log-domain code in commit 2. See Backend.logsumexp: keepdims missing from Cupy and Tensorflow backends; Cupy reimplements cupyx.scipy.special.logsumexp #867.
• Commit 2 generalizes convolutional_barycenter2d/convolutional_barycenter2d_debiased from hardcoded 2D images to arbitrary-dimension regular grids (1D signals, 3D volumes, ...), via a new _SeparableKernel applied through dense matmul (exp-domain) or a stabilized log-matmul-exp (log-domain), and vectorizes the log-domain solvers across histograms instead of looping with in-place writes. convolutional_barycenter2d{,_debiased} become thin A.ndim == 3 wrappers with unchanged signatures, defaults, docstrings and return contract; they are not deprecated.

Motivation

This is related to #862 (faster exact EMD on shared Cartesian grids). That discussion is about exact, network-flow-based OT on grids; this PR is a step toward the entropic/approximate side of the same problem space: generalizing the convolutional-kernel machinery to N-D grids is the natural prerequisite for adding the convolutional Wasserstein distance (Solomon et al. 2015, Algorithm 1) as a Sinkhorn-based alternative for OT/barycenters on d-dimensional grids — a follow-up PR builds on top of what's here. This PR does not implement that distance function itself.

Why matmul and not a native 1-D convolution

The Gaussian kernel on a grid axis is Toeplitz, so applying it via K @ x is mathematically equivalent to a (zero-padded) 1-D convolution — verified to agree to ~3e-15 on a small example — but a native convolution isn't used here, for two reasons:

1. Speed, in the regime POT cares about. Dense matmul goes through BLAS GEMM. Measured (Torch 2.14, CPU, float64, 5 images, n=512, reg=4e-3): ~4.5 ms for matmul vs. ~220 ms for a depthwise conv2d with a 4-sigma-truncated kernel.
2. Truncating the kernel is not a safe optimization for Sinkhorn. The dual scalings grow to compensate for the kernel's decay, so the transport plan stays spread out even where the kernel itself is negligible. On this package's own example (examples/barycenters/plot_convolutional_barycenter.py, 64x64 images, reg=0.004, so σ≈2.8 px), truncating the kernel below 12σ makes the Sinkhorn iteration diverge to NaN. Because reg is a physical quantity, σ in pixels grows with the grid for fixed reg (~2.8 px at n=64 vs. ~22.9 px at n=512), so truncation could only ever be a constant-factor saving, never an asymptotic one.

This is documented as a .. note:: on _SeparableKernel so it isn't relitigated in review.

Correctness verification

• Operator equivalence. Before deleting the old _get_convol_img_fn, the new _SeparableKernel was checked against it on random (4, 9, 7) input: exp-domain max abs diff ~1.8e-15; log-domain (fast, shifted) ~4.4e-16; log-domain against a module-private exact-logsumexp test helper (_exact_separable_log_apply, not part of the public API) exactly 0.0.
• Fast vs. exact log-domain numerics. The shifted log-matmul-exp used by _SeparableKernel in log-domain is a numerics change relative to the exact per-axis logsumexp reduction (not bit-identical). Regression-tested to agree to ~1e-9 (measured ~1e-15/1e-16) for reg >= 1e-3 on a moderate grid. Known, tested limitation: for an exact delta input at reg <= 1e-4 it underflows to -inf where the exact reduction stays finite.
• Jax/Tensorflow. method="sinkhorn_log" previously raised NotImplementedError for Jax/Tensorflow because of in-place writes (log_KU[k] = ...) in a per-histogram Python loop. Vectorizing across histograms removes the need for that, and both backends now genuinely pass (verified in a separate venv with tensorflow installed, since it isn't in this session's default env).
• Existing tests. All existing tests in test/test_bregman.py pass unchanged, except the jax/TF guards described above. No test tolerance needed adjusting.
• log=True return contract. _convolutional_barycenter2d puts U/V in the log dict; the _log variant does not. This asymmetry is preserved and now applied consistently across all four new N-D functions.

Benchmark

4 example images (examples/barycenters/plot_convolutional_barycenter.py), reg=0.004, numItermax=1000:

config	old (ms)	new (ms)	speedup	max|diff| (old vs new)
sinkhorn	228.6	31.2	7.3x	1.2e-18
sinkhorn_log	8591.0	129.3	66.5x	1.4e-18
sinkhorn (debiased)	344.1	46.5	7.4x	1.5e-17
sinkhorn_log (debiased)	12832.1	235.4	54.5x	2.7e-17

The large sinkhorn_log speedup comes from removing the per-histogram Python loop; the sinkhorn speedup comes from matmul-based batched kernel application replacing einsum. Old-vs-new agreement is at floating-point noise level.

Benchmark script (not part of this PR's tracked files — kept locally under the gitignored local_sandbox/)

# -*- coding: utf-8 -*-
"""
Benchmark for the convolutional_barycenter2d{,_debiased} refactor (PR: N-D
convolutional Wasserstein barycenters).

Compares the current (refactored) `ot.bregman.convolutional_barycenter2d`
against a self-contained copy of the pre-refactor implementation (the
per-histogram, `_get_convol_img_fn`-based version from
`ot/bregman/_convolutional.py` before this PR), on the 4-image example used
in `examples/barycenters/plot_convolutional_barycenter.py`.

Not part of the test suite: this is a manual benchmark, run with

    python local_sandbox/bench_convolutional_barycenter.py
"""

import time
import warnings

import numpy as np
import matplotlib.pyplot as plt

import ot

# --------------------------------------------------------------------------
# self-contained copy of the pre-refactor implementation, for comparison
# --------------------------------------------------------------------------


def _old_get_convol_img_fn(nx, width, height, reg, type_as, log_domain=False):
    t1 = nx.linspace(0, 1, width, type_as=type_as)
    Y1, X1 = nx.meshgrid(t1, t1)
    M1 = -((X1 - Y1) ** 2) / reg

    t2 = nx.linspace(0, 1, height, type_as=type_as)
    Y2, X2 = nx.meshgrid(t2, t2)
    M2 = -((X2 - Y2) ** 2) / reg

    if not log_domain:
        K1, K2 = nx.exp(M1), nx.exp(M2)

        def convol_imgs(imgs):
            kx = nx.einsum("...ij,kjl->kil", K1, imgs)
            kxy = nx.einsum("...ij,klj->kli", K2, kx)
            return kxy
    else:

        def convol_imgs(log_imgs):
            log_imgs = nx.logsumexp(M1[:, :, None] + log_imgs[None], axis=1)
            log_imgs = nx.logsumexp(M2[:, :, None] + log_imgs.T[None], axis=1).T
            return log_imgs

    return convol_imgs


def old_convolutional_barycenter2d(
    A, reg, weights=None, method="sinkhorn", numItermax=10000, stopThr=1e-4
):
    nx = ot.backend.get_backend(A)
    if method.lower() == "sinkhorn":
        return _old_convolutional_barycenter2d(
            A, reg, weights=weights, numItermax=numItermax, stopThr=stopThr
        )
    elif method.lower() == "sinkhorn_log":
        return _old_convolutional_barycenter2d_log(
            A, reg, weights=weights, numItermax=numItermax, stopThr=stopThr, nx=nx
        )
    raise ValueError(method)


def _old_convolutional_barycenter2d(
    A, reg, weights=None, numItermax=10000, stopThr=1e-9, stabThr=1e-30
):
    nx = ot.backend.get_backend(A)
    n_hists, width, height = A.shape
    if weights is None:
        weights = nx.ones((n_hists,), type_as=A) / n_hists

    bar = nx.ones((width, height), type_as=A)
    bar /= nx.sum(bar)
    U = nx.ones(A.shape, type_as=A)
    V = nx.ones(A.shape, type_as=A)
    convol_imgs = _old_get_convol_img_fn(nx, width, height, reg, type_as=A)

    KU = convol_imgs(U)
    for ii in range(numItermax):
        V = bar[None] / KU
        KV = convol_imgs(V)
        U = A / KV
        KU = convol_imgs(U)
        bar = nx.exp(nx.sum(weights[:, None, None] * nx.log(KU + stabThr), axis=0))
        if ii % 10 == 9:
            err = nx.sum(nx.std(V * KU, axis=0))
            if err < stopThr:
                break
    return bar


def _old_convolutional_barycenter2d_log(
    A, reg, weights=None, numItermax=10000, stopThr=1e-4, stabThr=1e-30, nx=None
):
    n_hists, width, height = A.shape
    if weights is None:
        weights = nx.ones((n_hists,), type_as=A) / n_hists

    convol_img = _old_get_convol_img_fn(nx, width, height, reg, type_as=A, log_domain=True)
    logA = nx.log(A + stabThr)
    log_KU, G = nx.zeros((2, *logA.shape), type_as=A)
    log_bar = nx.zeros((width, height), type_as=A)
    for ii in range(numItermax):
        log_bar = nx.zeros((width, height), type_as=A)
        for k in range(n_hists):
            f = logA[k] - convol_img(G[k])
            log_KU[k] = convol_img(f)
            log_bar = log_bar + weights[k] * log_KU[k]
        if ii % 10 == 9:
            err = nx.exp(G + log_KU).std(axis=0).sum()
            if err < stopThr:
                break
        G = log_bar[None, :, :] - log_KU
    return nx.exp(log_bar)


def old_convolutional_barycenter2d_debiased(
    A, reg, weights=None, method="sinkhorn", numItermax=10000, stopThr=1e-3
):
    nx = ot.backend.get_backend(A)
    if method.lower() == "sinkhorn":
        return _old_convolutional_barycenter2d_debiased(
            A, reg, weights=weights, numItermax=numItermax, stopThr=stopThr, nx=nx
        )
    elif method.lower() == "sinkhorn_log":
        return _old_convolutional_barycenter2d_debiased_log(
            A, reg, weights=weights, numItermax=numItermax, stopThr=stopThr, nx=nx
        )
    raise ValueError(method)


def _old_convolutional_barycenter2d_debiased(
    A, reg, weights=None, numItermax=10000, stopThr=1e-3, stabThr=1e-15, nx=None
):
    n_hists, width, height = A.shape
    if weights is None:
        weights = nx.ones((n_hists,), type_as=A) / n_hists

    bar = nx.ones((width, height), type_as=A)
    bar /= width * height
    U = nx.ones(A.shape, type_as=A)
    V = nx.ones(A.shape, type_as=A)
    c = nx.ones((width, height), type_as=A)
    convol_imgs = _old_get_convol_img_fn(nx, width, height, reg, type_as=A)

    KU = convol_imgs(U)
    for ii in range(numItermax):
        V = bar[None] / KU
        KV = convol_imgs(V)
        U = A / KV
        KU = convol_imgs(U)
        bar = c * nx.exp(nx.sum(weights[:, None, None] * nx.log(KU + stabThr), axis=0))
        for _ in range(10):
            c = (c * bar / nx.squeeze(convol_imgs(c[None]))) ** 0.5
        if ii % 10 == 9:
            err = nx.sum(nx.std(V * KU, axis=0))
            if err < stopThr and ii > 20:
                break
    return bar


def _old_convolutional_barycenter2d_debiased_log(
    A, reg, weights=None, numItermax=10000, stopThr=1e-3, stabThr=1e-30, nx=None
):
    n_hists, width, height = A.shape
    if weights is None:
        weights = nx.ones((n_hists,), type_as=A) / n_hists

    convol_img = _old_get_convol_img_fn(nx, width, height, reg, type_as=A, log_domain=True)
    logA = nx.log(A + stabThr)
    log_bar, c = nx.zeros((2, width, height), type_as=A)
    log_KU, G = nx.zeros((2, *logA.shape), type_as=A)
    for ii in range(numItermax):
        log_bar = nx.zeros((width, height), type_as=A)
        for k in range(n_hists):
            f = logA[k] - convol_img(G[k])
            log_KU[k] = convol_img(f)
            log_bar = log_bar + weights[k] * log_KU[k]
        log_bar = log_bar + c
        for _ in range(10):
            c = 0.5 * (c + log_bar - convol_img(c))
        if ii % 10 == 9:
            err = nx.sum(nx.std(nx.exp(G + log_KU), axis=0))
            if err < stopThr and ii > 20:
                break
        G = log_bar[None, :, :] - log_KU
    return nx.exp(log_bar)


# --------------------------------------------------------------------------
# benchmark
# --------------------------------------------------------------------------


def load_example_images():
    data_path = "data"
    names = ["redcross", "tooth", "heart", "duck"]
    imgs = [1 - plt.imread(f"{data_path}/{n}.png")[::2, ::2, 2] for n in names]
    imgs = [im / im.sum() for im in imgs]
    return np.array(imgs).astype(np.float64)


def bench(fn, *args, n_runs=3, **kwargs):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        fn(*args, **kwargs)  # warmup
        times = []
        for _ in range(n_runs):
            t0 = time.perf_counter()
            out = fn(*args, **kwargs)
            times.append(time.perf_counter() - t0)
    return min(times), out


def main():
    A = load_example_images()
    print(f"A.shape = {A.shape}")
    reg = 0.004

    configs = [
        ("sinkhorn", ot.bregman.convolutional_barycenter2d, old_convolutional_barycenter2d, 1e-4),
        (
            "sinkhorn_log",
            ot.bregman.convolutional_barycenter2d,
            old_convolutional_barycenter2d,
            1e-4,
        ),
        (
            "sinkhorn (debiased)",
            ot.bregman.convolutional_barycenter2d_debiased,
            old_convolutional_barycenter2d_debiased,
            1e-3,
        ),
        (
            "sinkhorn_log (debiased)",
            ot.bregman.convolutional_barycenter2d_debiased,
            old_convolutional_barycenter2d_debiased,
            1e-3,
        ),
    ]

    print(f"{'config':<26}{'old (ms)':>12}{'new (ms)':>12}{'speedup':>10}{'max|diff|':>14}")
    for label, new_fn, old_fn, stop_thr in configs:
        method = "sinkhorn_log" if "sinkhorn_log" in label else "sinkhorn"
        t_old, bar_old = bench(
            old_fn, A, reg, method=method, numItermax=1000, stopThr=stop_thr
        )
        t_new, bar_new = bench(
            new_fn, A, reg, method=method, numItermax=1000, stopThr=stop_thr
        )
        diff = np.max(np.abs(np.asarray(bar_old) - np.asarray(bar_new)))
        print(
            f"{label:<26}{t_old * 1000:>12.2f}{t_new * 1000:>12.2f}"
            f"{t_old / t_new:>9.2f}x{diff:>14.3e}"
        )


if __name__ == "__main__":
    main()

New public API

• ot.bregman.convolutional_grid_barycenter(A, reg, ...) — A of shape (n_hists, *grid_shape), any number of grid dimensions.
• ot.bregman.convolutional_grid_barycenter_debiased(A, reg, ...) — debiased variant.

convolutional_barycenter2d/convolutional_barycenter2d_debiased are unchanged thin wrappers (A.ndim == 3 check, then delegate).

Tests

• New: 1D ((3, 40)) and 3D ((3, 12, 13, 14)) barycenters, both methods, debiased and not, parametrized over backends.
• New: convolutional_barycenter2d{,_debiased} vs convolutional_grid_barycenter{,_debiased} equivalence for A.ndim == 3, and ValueError on wrong ndim.
• New: fast-vs-exact log-domain regression test, and a characterization test of the documented underflow limitation.
• Updated: jax/TF sinkhorn_log guards in test/test_bregman.py now exercise the real solver instead of expecting NotImplementedError.
• Also extended test/test_backend.py for the Backend.logsumexp fix (axis, keepdims, all--inf row).

Out of scope / deferred to follow-ups

• The convolutional Wasserstein distance / Sinkhorn divergence on grids (next PR, motivated by Faster exact EMD for distributions on shared multi-dimensional cartesian grids (with cityblock or sqeuclidean ground metric) #862).
• Kernel truncation, FFT convolution, IIR/recursive Gaussian filtering, or a nx.convolve1d backend primitive (see "Why matmul" above).
• An anisotropic spacing/extent argument for non-unit grids, or any change to the linspace(0, 1, n) grid convention.
• Deprecating the 2D entry points.

🤖 Generated with Claude Code