API REFERENCE/omnibias-graph

omnibias.graph.torch.ops.relaxation

0.1.0a2AlphaPython >=3.10AGPL / commercial

Spectral graph operators and differentiable relaxations for ordering, assignment, and graph-structured learning.

import omnibias.graph.torch.ops.relaxation

Module contract & mathematical conventions

Differentiable relaxations of discrete combinatorial objects (torch).

Each operator carries a temperature tau > 0 that recovers the hard object as tau -> 0:

• sinkhorn_normalize -- projection onto the Birkhoff polytope (doubly-stochastic matrices) by log-domain matrix scaling, mirroring the Sinkhorn matrix scaling.
• gumbel_sinkhorn -- a differentiable relaxation of the assignment / matching problem (Mena et al., 2018): Sinkhorn-normalise (log_alpha + noise) / tau.
• soft_sort / soft_sort_permutation -- the SoftSort operator (Prillo & Eisenschlos, 2020): a row-stochastic soft permutation that becomes the exact argsort permutation as tau -> 0.
• soft_top_k -- membership weights in [0, 1] summing to exactly k, built from the SoftSort permutation.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

FUNCTION

gumbel_sinkhorn

omnibias.graph.torch.ops.gumbel_sinkhornSource ↗

gumbel_sinkhorn(log_alpha: Tensor, *, temperature: float=1.0, n_iters: int=20, noise: Tensor | None=None) -> Tensor

Gumbel-Sinkhorn relaxation of a permutation / assignment matrix.

Returns sinkhorn_normalize((log_alpha + noise) / temperature). Pass a pre-sampled Gumbel noise tensor (same shape as log_alpha) for the stochastic operator; with noise=None the map is deterministic. As temperature -> 0 the result approaches a hard permutation matrix.

Parameters

Name	Type & default	Contract
log_alphapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
temperaturekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
n_iterskeyword-only	intDefault: 20	See the operation contract and implementation below.
noisekeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When temperature <= 0.0, raises ValueError(f'temperature must be > 0; got {temperature}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · temperature must be positive

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_temperature_must_be_positive(self) -> None:
        with pytest.raises(ValueError, match="temperature"):
            G.gumbel_sinkhorn(torch.eye(3), temperature=0.0)

python -m pytest "packages/omnibias-graph/tests/test_relaxation_torch.py::TestGumbelSinkhorn::test_temperature_must_be_positive" -q

Regression example · low temperature is permutation

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_low_temperature_is_permutation(self) -> None:
        # A near-identity affinity: at low temperature Gumbel-Sinkhorn (no noise)
        # recovers the identity permutation matrix.
        log_alpha = 3.0 * torch.eye(5)
        p = G.gumbel_sinkhorn(log_alpha, temperature=0.01, n_iters=200)
        assert torch.allclose(p, torch.eye(5), atol=1e-3)

python -m pytest packages/omnibias-graph/tests/test_relaxation_torch.py -k test_low_temperature_is_permutation -q

Read the implementation

def gumbel_sinkhorn(
    log_alpha: Tensor,
    *,
    temperature: float = 1.0,
    n_iters: int = 20,
    noise: Tensor | None = None,
) -> Tensor:
    r"""Gumbel-Sinkhorn relaxation of a permutation / assignment matrix.

    Returns ``sinkhorn_normalize((log_alpha + noise) / temperature)``. Pass a
    pre-sampled Gumbel ``noise`` tensor (same shape as ``log_alpha``) for the
    stochastic operator; with ``noise=None`` the map is deterministic. As
    ``temperature -> 0`` the result approaches a hard permutation matrix.
    """
    if temperature <= 0.0:
        raise ValueError(f"temperature must be > 0; got {temperature}")
    perturbed = log_alpha if noise is None else log_alpha + noise
    return sinkhorn_normalize(perturbed / temperature, n_iters=n_iters)

Related APIsomnibias.graph.torch.ops.sinkhorn_normalize↗

Other public import paths

• omnibias.graph.torch.ops.gumbel_sinkhorn
• omnibias.graph.torch.ops.relaxation.gumbel_sinkhorn

FUNCTION

sample_gumbel

omnibias.graph.torch.ops.sample_gumbelSource ↗

sample_gumbel(shape: tuple[int, ...], *, dtype: torch.dtype | None=None, device: torch.device | None=None, generator: torch.Generator | None=None, eps: float=1e-20) -> Tensor

I.i.d. standard Gumbel noise -log(-log(U)), U ~ Uniform(0, 1).

Parameters

Name	Type & default	Contract
shapepositional-or-keyword	tuple[int, ...]Required	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.
devicekeyword-only	torch.device | NoneDefault: None	See the operation contract and implementation below.
generatorkeyword-only	torch.Generator | NoneDefault: None	See the operation contract and implementation below.
epskeyword-only	floatDefault: 1e-20	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def sample_gumbel(
    shape: tuple[int, ...],
    *,
    dtype: torch.dtype | None = None,
    device: torch.device | None = None,
    generator: torch.Generator | None = None,
    eps: float = 1e-20,
) -> Tensor:
    r"""I.i.d. standard Gumbel noise ``-log(-log(U))``, ``U ~ Uniform(0, 1)``."""
    u = torch.rand(shape, dtype=dtype, device=device, generator=generator)
    return -torch.log(-torch.log(u + eps) + eps)

Other public import paths

• omnibias.graph.torch.ops.sample_gumbel
• omnibias.graph.torch.ops.relaxation.sample_gumbel

FUNCTION

sinkhorn_normalize

omnibias.graph.torch.ops.sinkhorn_normalizeSource ↗

sinkhorn_normalize(log_alpha: Tensor, *, n_iters: int=20) -> Tensor

Project exp(log_alpha) onto the doubly-stochastic matrices.

log_alpha is a square (n, n) log-domain matrix. Alternating row / column log-normalisation converges to the unique doubly-stochastic matrix diag(u) exp(log_alpha) diag(v) (Sinkhorn's theorem). Returns the normalised matrix in the probability domain.

Parameters

Name	Type & default	Contract
log_alphapositional-or-keyword	TensorRequired	See the operation contract and implementation below.
n_iterskeyword-only	intDefault: 20	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When log_alpha.ndim != 2 or log_alpha.shape[-1] != log_alpha.shape[-2], raises ValueError(f'log_alpha must be square (n, n); got {tuple(log_alpha.shape)}').
• When n_iters < 1, raises ValueError(f'n_iters must be >= 1; got {n_iters}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · non square raises

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_non_square_raises(self) -> None:
        with pytest.raises(ValueError, match="square"):
            G.sinkhorn_normalize(torch.zeros(3, 4))

python -m pytest "packages/omnibias-graph/tests/test_relaxation_torch.py::TestSinkhorn::test_non_square_raises" -q

Regression example · sinkhorn parity

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_sinkhorn_parity() -> None:
    rng = np.random.default_rng(2)
    la = rng.normal(size=(5, 5))
    _assert_close(
        GT.sinkhorn_normalize(torch.tensor(la), n_iters=100),
        GJ.sinkhorn_normalize(jnp.asarray(la), n_iters=100),
    )

python -m pytest packages/omnibias-graph/tests/test_cross_backend.py -k test_sinkhorn_parity -q

Read the implementation

def sinkhorn_normalize(log_alpha: Tensor, *, n_iters: int = 20) -> Tensor:
    r"""Project ``exp(log_alpha)`` onto the doubly-stochastic matrices.

    ``log_alpha`` is a square ``(n, n)`` log-domain matrix. Alternating row / column
    log-normalisation converges to the unique doubly-stochastic matrix
    ``diag(u) exp(log_alpha) diag(v)`` (Sinkhorn's theorem). Returns the
    normalised matrix in the probability domain.
    """
    if log_alpha.ndim != 2 or log_alpha.shape[-1] != log_alpha.shape[-2]:
        raise ValueError(
            f"log_alpha must be square (n, n); got {tuple(log_alpha.shape)}"
        )
    if n_iters < 1:
        raise ValueError(f"n_iters must be >= 1; got {n_iters}")
    log_p = log_alpha
    for _ in range(n_iters):
        log_p = log_p - torch.logsumexp(log_p, dim=-1, keepdim=True)
        log_p = log_p - torch.logsumexp(log_p, dim=-2, keepdim=True)
    return torch.exp(log_p)

Other public import paths

• omnibias.graph.torch.ops.sinkhorn_normalize
• omnibias.graph.torch.ops.relaxation.sinkhorn_normalize

FUNCTION

soft_sort

omnibias.graph.torch.ops.soft_sortSource ↗

soft_sort(scores: Tensor, *, temperature: float=1.0, descending: bool=True) -> Tensor

Differentiable sorted values P @ scores -> (n,).

Approaches torch.sort(scores) as temperature -> 0.

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
temperaturekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
descendingkeyword-only	boolDefault: True	See the operation contract and implementation below.

ReturnsTensor

Executed regression · differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_differentiable(self) -> None:
        s = torch.tensor([3.0, 1.0, 4.0, 1.5], requires_grad=True)
        G.soft_sort(s, temperature=0.7).sum().backward()
        assert s.grad is not None and torch.isfinite(s.grad).all()

python -m pytest "packages/omnibias-graph/tests/test_relaxation_torch.py::TestSoftSort::test_differentiable" -q

Regression example · soft sort parity

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_soft_sort_parity(tau: float) -> None:
    rng = np.random.default_rng(4)
    s = rng.normal(size=7)
    _assert_close(
        GT.soft_sort(torch.tensor(s), temperature=tau),
        GJ.soft_sort(jnp.asarray(s), temperature=tau),
    )

python -m pytest packages/omnibias-graph/tests/test_cross_backend.py -k test_soft_sort_parity -q

Read the implementation

def soft_sort(
    scores: Tensor, *, temperature: float = 1.0, descending: bool = True
) -> Tensor:
    r"""Differentiable sorted values ``P @ scores`` -> ``(n,)``.

    Approaches ``torch.sort(scores)`` as ``temperature -> 0``.
    """
    p = soft_sort_permutation(scores, temperature=temperature, descending=descending)
    return p @ scores.reshape(-1)

Related APIsomnibias.graph.torch.ops.soft_sort_permutation↗

Other public import paths

• omnibias.graph.torch.ops.soft_sort
• omnibias.graph.torch.ops.relaxation.soft_sort

FUNCTION

soft_sort_permutation

omnibias.graph.torch.ops.soft_sort_permutationSource ↗

soft_sort_permutation(scores: Tensor, *, temperature: float=1.0, descending: bool=True) -> Tensor

SoftSort row-stochastic soft permutation for a score vector (n,).

P[i, j] = softmax_j(-|sorted(scores)[i] - scores[j]| / tau); row i is a soft one-hot selecting the input that is the i-th in sorted order. As tau -> 0 (distinct scores) P becomes the exact argsort permutation.

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
temperaturekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.
descendingkeyword-only	boolDefault: True	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When temperature <= 0.0, raises ValueError(f'temperature must be > 0; got {temperature}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · differentiable

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_differentiable(self) -> None:
        s = torch.tensor([3.0, 1.0, 4.0, 1.5], requires_grad=True)
        G.soft_sort(s, temperature=0.7).sum().backward()
        assert s.grad is not None and torch.isfinite(s.grad).all()

python -m pytest "packages/omnibias-graph/tests/test_relaxation_torch.py::TestSoftSort::test_differentiable" -q

Regression example · permutation row stochastic

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_permutation_row_stochastic(self) -> None:
        s = torch.tensor([3.0, 1.0, 4.0, 1.5, 2.0])
        p = G.soft_sort_permutation(s, temperature=0.5)
        assert torch.allclose(p.sum(dim=-1), torch.ones(5), atol=1e-10)
        assert (p >= 0).all()

python -m pytest packages/omnibias-graph/tests/test_relaxation_torch.py -k test_permutation_row_stochastic -q

Read the implementation

def soft_sort_permutation(
    scores: Tensor, *, temperature: float = 1.0, descending: bool = True
) -> Tensor:
    r"""SoftSort row-stochastic soft permutation for a score vector ``(n,)``.

    ``P[i, j] = softmax_j(-|sorted(scores)[i] - scores[j]| / tau)``; row ``i`` is a
    soft one-hot selecting the input that is the ``i``-th in sorted order. As
    ``tau -> 0`` (distinct scores) ``P`` becomes the exact argsort permutation.
    """
    if temperature <= 0.0:
        raise ValueError(f"temperature must be > 0; got {temperature}")
    s = scores.reshape(-1)
    sorted_s = torch.sort(s, descending=descending).values
    dist = torch.abs(sorted_s.unsqueeze(-1) - s.unsqueeze(-2))
    return torch.softmax(-dist / temperature, dim=-1)

Other public import paths

• omnibias.graph.torch.ops.soft_sort_permutation
• omnibias.graph.torch.ops.relaxation.soft_sort_permutation

FUNCTION

soft_top_k

omnibias.graph.torch.ops.soft_top_kSource ↗

soft_top_k(scores: Tensor, k: int, *, temperature: float=1.0) -> Tensor

Soft top-k membership weights in [0, 1] summing to exactly k.

m_j = sum_{i < k} P[i, j] with P the descending SoftSort permutation: the (relaxed) probability that input j is among the k largest scores. As temperature -> 0 this is the hard top-k indicator.

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.
temperaturekeyword-only	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When k < 1 or k > n, raises ValueError(f'k must be in [1, {n}]; got {k}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · bad k raises

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

    def test_bad_k_raises(self) -> None:
        with pytest.raises(ValueError, match=r"k must be"):
            G.soft_top_k(torch.zeros(4), 5)

python -m pytest "packages/omnibias-graph/tests/test_relaxation_torch.py::TestSoftTopK::test_bad_k_raises" -q

Regression example · soft top k parity

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_soft_top_k_parity() -> None:
    rng = np.random.default_rng(11)
    s = rng.normal(size=8)
    _assert_close(
        GT.soft_top_k(torch.tensor(s), 3, temperature=0.2),
        GJ.soft_top_k(jnp.asarray(s), 3, temperature=0.2),
    )

python -m pytest packages/omnibias-graph/tests/test_cross_backend.py -k test_soft_top_k_parity -q

Read the implementation

def soft_top_k(
    scores: Tensor, k: int, *, temperature: float = 1.0
) -> Tensor:
    r"""Soft top-``k`` membership weights in ``[0, 1]`` summing to exactly ``k``.

    ``m_j = sum_{i < k} P[i, j]`` with ``P`` the descending SoftSort permutation:
    the (relaxed) probability that input ``j`` is among the ``k`` largest scores.
    As ``temperature -> 0`` this is the hard top-``k`` indicator.
    """
    n = scores.reshape(-1).shape[0]
    if k < 1 or k > n:
        raise ValueError(f"k must be in [1, {n}]; got {k}")
    p = soft_sort_permutation(scores, temperature=temperature, descending=True)
    return p[:k, :].sum(dim=0)

Related APIsomnibias.graph.torch.ops.soft_sort_permutation↗

Other public import paths

• omnibias.graph.torch.ops.soft_top_k
• omnibias.graph.torch.ops.relaxation.soft_top_k

← Package overviewRunnable learning examples →