waxmorph.jax.losses.make_sinkhorn_loss

waxmorph.jax.losses.make_sinkhorn_loss#

waxmorph.jax.losses.make_sinkhorn_loss(blur=0.05, *, p=2, cost_fn=None, **solve_kwargs)[source]#

Build OTT’s debiased entropic-OT divergence for unordered clouds.

\[S_\varepsilon(\alpha,\beta)=\mathrm{OT}_\varepsilon(\alpha,\beta) -\tfrac12\mathrm{OT}_\varepsilon(\alpha,\alpha) -\tfrac12\mathrm{OT}_\varepsilon(\beta,\beta).\]

blur sets epsilon = blur ** p. With cost_fn=None, p=1 selects Euclidean cost and p=2 selects half squared Euclidean cost; other exponents use an explicit cost. An explicit cost replaces that selection, while p still sets the epsilon exponent. solve_kwargs are forwarded to OTT.

The factory provides debiased Sinkhorn divergence with OTT options. PyTorch’s GeomLoss factory additionally provides MMD, Hausdorff, scaling, and reach contracts. Numerical values follow each solver’s options. OTT is imported lazily.

Raises:

ImportError – If importing ott-jax fails.

Parameters:
Return type:

Callable[[Array, Array], Array]