API REFERENCE/omnibias-graph

omnibias.graph.jax.ops.relaxation

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.graph.jax.ops.relaxation

Module contract & mathematical conventions

Differentiable relaxations of discrete combinatorial objects (jax).

Bit-identical twin of omnibias.graph.torch.ops.relaxation; see that module for the definitions and references. Each operator carries a temperature tau > 0 that recovers the hard object as tau -> 0.

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 →

External type aliases

These public names refer to external types. Their implementations and inherited members belong to the originating library.

Public name	External definition
omnibias.graph.jax.ops.relaxation.Array	typing.Any

FUNCTION

gumbel_sinkhorn

omnibias.graph.jax.ops.gumbel_sinkhornSource ↗

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

Gumbel-Sinkhorn relaxation of a permutation / assignment matrix.

Parameters

Name	Type & default	Contract
log_alphapositional-or-keyword	ArrayRequired	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	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

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 · gumbel sinkhorn recovers permutation

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_gumbel_sinkhorn_recovers_permutation() -> None:
    perm = np.array([4, 3, 2, 1, 0])
    log_alpha = np.full((5, 5), -5.0)
    for i, j in enumerate(perm):
        log_alpha[i, j] = 5.0
    p = np.asarray(G.gumbel_sinkhorn(jnp.asarray(log_alpha), temperature=0.05, n_iters=200))
    assert np.array_equal(p.argmax(axis=-1), perm)

python -m pytest "packages/omnibias-graph/tests/test_relaxation_jax.py::test_gumbel_sinkhorn_recovers_permutation" -q

Regression example · gumbel sinkhorn parity with shared noise

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_gumbel_sinkhorn_parity_with_shared_noise() -> None:
    rng = np.random.default_rng(9)
    la = rng.normal(size=(5, 5))
    noise = rng.gumbel(size=(5, 5))
    _assert_close(
        GT.gumbel_sinkhorn(torch.tensor(la), temperature=0.3, n_iters=100, noise=torch.tensor(noise)),
        GJ.gumbel_sinkhorn(jnp.asarray(la), temperature=0.3, n_iters=100, noise=jnp.asarray(noise)),
    )

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

Read the implementation

def gumbel_sinkhorn(
    log_alpha: Array,
    *,
    temperature: float = 1.0,
    n_iters: int = 20,
    noise: Array | None = None,
) -> Array:
    r"""Gumbel-Sinkhorn relaxation of a permutation / assignment 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.jax.ops.sinkhorn_normalize↗

Other public import paths

• omnibias.graph.jax.ops.gumbel_sinkhorn
• omnibias.graph.jax.ops.relaxation.gumbel_sinkhorn

FUNCTION

sample_gumbel

omnibias.graph.jax.ops.sample_gumbelSource ↗

sample_gumbel(shape: tuple[int, ...], key: Any, *, dtype: Any=None, eps: float=1e-20) -> Array

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.
keypositional-or-keyword	AnyRequired	See the operation contract and implementation below.
dtypekeyword-only	AnyDefault: None	See the operation contract and implementation below.
epskeyword-only	floatDefault: 1e-20	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def sample_gumbel(
    shape: tuple[int, ...],
    key: Any,
    *,
    dtype: Any = None,
    eps: float = 1e-20,
) -> Array:
    r"""I.i.d. standard Gumbel noise ``-log(-log(U))``, ``U ~ Uniform(0, 1)``."""
    u = jax.random.uniform(key, shape, dtype=dtype or jnp.float64)
    return -jnp.log(-jnp.log(u + eps) + eps)

Other public import paths

• omnibias.graph.jax.ops.sample_gumbel
• omnibias.graph.jax.ops.relaxation.sample_gumbel

FUNCTION

sinkhorn_normalize

omnibias.graph.jax.ops.sinkhorn_normalizeSource ↗

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

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

Parameters

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

ReturnsArray

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 · sinkhorn parity

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_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::test_sinkhorn_parity" -q

Regression example · sinkhorn doubly 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_sinkhorn_doubly_stochastic() -> None:
    rng = np.random.default_rng(0)
    p = np.asarray(G.sinkhorn_normalize(jnp.asarray(rng.normal(size=(6, 6))), n_iters=300))
    assert np.allclose(p.sum(axis=-1), 1.0, atol=1e-9)
    assert np.allclose(p.sum(axis=-2), 1.0, atol=1e-9)

python -m pytest packages/omnibias-graph/tests/test_relaxation_jax.py -k test_sinkhorn_doubly_stochastic -q

Read the implementation

def sinkhorn_normalize(log_alpha: Array, *, n_iters: int = 20) -> Array:
    r"""Project ``exp(log_alpha)`` onto the doubly-stochastic matrices."""
    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 - logsumexp(log_p, axis=-1, keepdims=True)
        log_p = log_p - logsumexp(log_p, axis=-2, keepdims=True)
    return jnp.exp(log_p)

Other public import paths

• omnibias.graph.jax.ops.sinkhorn_normalize
• omnibias.graph.jax.ops.relaxation.sinkhorn_normalize

FUNCTION

soft_sort

omnibias.graph.jax.ops.soft_sortSource ↗

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

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

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	ArrayRequired	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.

ReturnsArray

Executed regression · temperature validation

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_validation() -> None:
    with pytest.raises(ValueError, match="temperature"):
        G.soft_sort(jnp.zeros(3), temperature=0.0)

python -m pytest "packages/omnibias-graph/tests/test_relaxation_jax.py::test_temperature_validation" -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: Array, *, temperature: float = 1.0, descending: bool = True
) -> Array:
    r"""Differentiable sorted values ``P @ scores`` -> ``(n,)``."""
    p = soft_sort_permutation(scores, temperature=temperature, descending=descending)
    return p @ scores.reshape(-1)

Related APIsomnibias.graph.jax.ops.soft_sort_permutation↗

Other public import paths

• omnibias.graph.jax.ops.soft_sort
• omnibias.graph.jax.ops.relaxation.soft_sort

FUNCTION

soft_sort_permutation

omnibias.graph.jax.ops.soft_sort_permutationSource ↗

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

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

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	ArrayRequired	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.

ReturnsArray

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 validation

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_validation() -> None:
    with pytest.raises(ValueError, match="temperature"):
        G.soft_sort(jnp.zeros(3), temperature=0.0)

python -m pytest "packages/omnibias-graph/tests/test_relaxation_jax.py::test_temperature_validation" -q

Read the implementation

def soft_sort_permutation(
    scores: Array, *, temperature: float = 1.0, descending: bool = True
) -> Array:
    r"""SoftSort row-stochastic soft permutation for a score vector ``(n,)``."""
    if temperature <= 0.0:
        raise ValueError(f"temperature must be > 0; got {temperature}")
    s = scores.reshape(-1)
    sorted_s = jnp.sort(s)
    if descending:
        sorted_s = sorted_s[::-1]
    dist = jnp.abs(sorted_s[:, None] - s[None, :])
    return jax.nn.softmax(-dist / temperature, axis=-1)

Other public import paths

• omnibias.graph.jax.ops.soft_sort_permutation
• omnibias.graph.jax.ops.relaxation.soft_sort_permutation

FUNCTION

soft_top_k

omnibias.graph.jax.ops.soft_top_kSource ↗

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

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

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	ArrayRequired	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.

ReturnsArray

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 · soft top k parity

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_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::test_soft_top_k_parity" -q

Regression example · soft top k sums to k and hard limit

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_sums_to_k_and_hard_limit() -> None:
    s = jnp.asarray([3.0, 1.0, 4.0, 1.5, 9.0, 2.0])
    for tau in (0.01, 0.5, 3.0):
        m = np.asarray(G.soft_top_k(s, 3, temperature=tau))
        assert abs(m.sum() - 3.0) < 1e-9
    m = np.asarray(G.soft_top_k(s, 3, temperature=1e-4))
    assert set(np.nonzero(m > 0.5)[0].tolist()) == {4, 2, 0}

python -m pytest packages/omnibias-graph/tests/test_relaxation_jax.py -k test_soft_top_k_sums_to_k_and_hard_limit -q

Read the implementation

def soft_top_k(scores: Array, k: int, *, temperature: float = 1.0) -> Array:
    r"""Soft top-``k`` membership weights in ``[0, 1]`` summing to exactly ``k``."""
    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(axis=0)

Related APIsomnibias.graph.jax.ops.soft_sort_permutation↗

Other public import paths

• omnibias.graph.jax.ops.soft_top_k
• omnibias.graph.jax.ops.relaxation.soft_top_k

← Package overviewRunnable learning examples →