API REFERENCE/omnibias-struct

omnibias.struct.torch.distributions

0.1.0a2AlphaPython >=3.10AGPL / commercial

Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.

import omnibias.struct.torch.distributions

Module contract & mathematical conventions

Distribution operators over the Gibbs path distribution p_beta (torch).

Bit-identical twin of omnibias.struct.jax.distributions for the differentiable ops. For the Gibbs distribution p_beta(D) proportional to exp(beta * score(D)) over the derivations of a ~omnibias.struct._core.semiring.Hypergraph:

• path_entropy -- the exact Shannon entropy H(p_beta) = beta (V_beta - E_p[score]) in closed form from the inside value and the edge marginals (the first-order expectation-semiring quantity); differentiable, >= 0, -> log(#derivations) as beta -> 0 and -> log(#argmax) as beta -> inf.
• sample_paths -- **exact** forward-filtering backward-sampling from p_beta (not a relaxation); its empirical edge frequencies converge to semiring_marginals.
• gumbel_relaxed_sample -- a Gumbel-perturb-and-MAP **approximate**, differentiable relaxed sample (optionally straight-through); honestly labelled, *not* an exact draw.
• topk_paths -- **exact** k-best derivations (a hard decode) and topk_free_energy -- the differentiable lse_beta restricted to those top-k scores (-> best score as beta -> inf, -> the full soft value as k -> #derivations).

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_relaxed_sample

omnibias.struct.torch.gumbel_relaxed_sampleSource ↗

gumbel_relaxed_sample(graph: Hypergraph, edge_weights: Tensor, beta: float=1.0, *, seed: int | None=None, hard: bool=False) -> Tensor

Gumbel perturb-and-MAP **approximate** relaxed sample (differentiable), as edge scores.

Adds i.i.d. Gumbel noise to the edge weights and returns the tower-softmax edge selection semiring_marginals of the perturbed weights; with hard=True a straight-through estimator returns the argmax-derivation indicator in the forward pass and the relaxed gradient in the backward. This is a per-edge perturbation, **not** an exact draw from p_beta (use sample_paths for that) -- honestly labelled as a relaxation.

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
seedkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.
hardkeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsTensor

Executed regression · gumbel relaxed sample

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_relaxed_sample() -> None:
    w = torch.tensor(_weights(0), requires_grad=True)
    beta = 4.0
    relaxed = st.gumbel_relaxed_sample(GRAPH, w, beta, seed=0)
    assert relaxed.shape == w.shape
    (g,) = torch.autograd.grad(relaxed.sum(), w)
    assert torch.isfinite(g).all()  # differentiable relaxed sample
    hard = st.gumbel_relaxed_sample(GRAPH, w.detach(), beta, seed=0, hard=True)
    # the hard straight-through forward value is an integer edge-count indicator of one derivation
    hard_np = hard.detach().numpy()
    assert hard.sum() > 0.0
    assert np.allclose(hard_np, np.round(hard_np), atol=1e-9)

python -m pytest "packages/omnibias-struct/tests/test_operators.py::test_gumbel_relaxed_sample" -q

Read the implementation

def gumbel_relaxed_sample(
    graph: Hypergraph,
    edge_weights: Tensor,
    beta: float = 1.0,
    *,
    seed: int | None = None,
    hard: bool = False,
) -> Tensor:
    r"""Gumbel perturb-and-MAP **approximate** relaxed sample (differentiable), as edge scores.

    Adds i.i.d. Gumbel noise to the edge weights and returns the tower-softmax edge selection
    :func:`semiring_marginals` of the perturbed weights; with ``hard=True`` a straight-through
    estimator returns the argmax-derivation indicator in the forward pass and the relaxed
    gradient in the backward. This is a per-edge perturbation, **not** an exact draw from
    ``p_beta`` (use :func:`sample_paths` for that) -- honestly labelled as a relaxation.
    """
    generator = torch.Generator(device=edge_weights.device)
    if seed is not None:
        generator.manual_seed(seed)
    u = torch.rand(edge_weights.shape, dtype=edge_weights.dtype, device=edge_weights.device, generator=generator)
    gumbel = -torch.log(-torch.log(u))
    perturbed = edge_weights + gumbel.detach()
    relaxed: Tensor = semiring_marginals(graph, perturbed, beta)
    if not hard:
        return relaxed
    _score, deriv = best_derivation(graph, [float(x) for x in perturbed.detach().reshape(-1).tolist()])
    hard_ind = torch.zeros_like(edge_weights)
    for e in deriv:
        hard_ind[e] = hard_ind[e] + 1.0
    straight_through: Tensor = hard_ind + (relaxed - relaxed.detach())
    return straight_through

Related APIsomnibias.struct.best_derivation↗omnibias.struct.torch.semiring_marginals↗

Other public import paths

• omnibias.struct.torch.gumbel_relaxed_sample
• omnibias.struct.torch.distributions.gumbel_relaxed_sample

FUNCTION

path_entropy

omnibias.struct.torch.path_entropySource ↗

path_entropy(graph: Hypergraph, edge_weights: Tensor, beta: float=1.0) -> Tensor

Exact Shannon entropy (nats) of p_beta over the derivations of graph.

Closed form H = beta (V_beta - E_p[score]) = beta (value - <marginals, weights>) -- the first-order expectation-semiring entropy. Differentiable in edge_weights; equals omnibias.struct._core.operators.brute_force_entropy on tiny graphs.

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · entropy identity and mc

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_entropy_identity_and_mc(backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_entropy(0, BETA, backend)
    assert rec.ok, rec
    assert rec.metrics["identity_err"] < 1e-8

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_entropy_identity_and_mc[torch]" -q

Regression example · path entropy is differentiable

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_path_entropy_is_differentiable(seed: int) -> None:
    w = torch.tensor(_weights(seed), requires_grad=True)
    h = st.path_entropy(GRAPH, w, 3.0)
    (g,) = torch.autograd.grad(h, w)
    assert torch.isfinite(g).all()

python -m pytest packages/omnibias-struct/tests/test_operators.py -k test_path_entropy_is_differentiable -q

Read the implementation

def path_entropy(graph: Hypergraph, edge_weights: Tensor, beta: float = 1.0) -> Tensor:
    r"""Exact Shannon entropy (nats) of ``p_beta`` over the derivations of ``graph``.

    Closed form ``H = beta (V_beta - E_p[score]) = beta (value - <marginals, weights>)`` --
    the first-order expectation-semiring entropy. Differentiable in ``edge_weights``; equals
    :func:`omnibias.struct._core.operators.brute_force_entropy` on tiny graphs.
    """
    value = semiring_value(graph, edge_weights, beta)
    marginals = semiring_marginals(graph, edge_weights, beta)
    expected_score = (marginals * edge_weights).sum()
    entropy: Tensor = beta * (value - expected_score)
    return entropy

Related APIsomnibias.struct.torch.semiring_marginals↗omnibias.struct.torch.semiring_value↗

Other public import paths

• omnibias.struct.torch.path_entropy
• omnibias.struct.torch.distributions.path_entropy

FUNCTION

sample_paths

omnibias.struct.torch.sample_pathsSource ↗

sample_paths(graph: Hypergraph, edge_weights: Tensor, beta: float=1.0, num_samples: int=1, *, seed: int | None=None) -> tuple[Tensor, list[tuple[int, ...]]]

Exact forward-filtering backward-sampling of num_samples derivations from p_beta.

Returns (counts, samples): counts is a (num_samples, num_edges) tensor of per-edge usage counts and samples the chosen edge-index tuples. Exact (not a relaxation); counts.mean(0) -> semiring_marginals as num_samples -> inf. Sampling is not differentiable (use gumbel_relaxed_sample for a relaxed draw).

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
num_samplespositional-or-keyword	intDefault: 1	See the operation contract and implementation below.
seedkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[Tensor, list[tuple[int, ...]]]

Executed regression · sampler empirical marginals converge

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_sampler_empirical_marginals_converge(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_sampling(0, BETA, backend).ok

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_sampler_empirical_marginals_converge[torch]" -q

Regression example · sample paths empirical marginals

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_sample_paths_empirical_marginals() -> None:
    w = _weights(0)
    wt = torch.tensor(w)
    beta = 2.0
    counts, samples = st.sample_paths(GRAPH, wt, beta, 6000, seed=0)
    assert len(samples) == 6000
    empirical = counts.mean(0).numpy()
    closed_form = st.semiring_marginals(GRAPH, wt, beta).numpy()
    assert float(np.max(np.abs(empirical - closed_form))) < 0.03  # Monte-Carlo error ~ 1/sqrt(6000)

python -m pytest packages/omnibias-struct/tests/test_operators.py -k test_sample_paths_empirical_marginals -q

Read the implementation

def sample_paths(
    graph: Hypergraph,
    edge_weights: Tensor,
    beta: float = 1.0,
    num_samples: int = 1,
    *,
    seed: int | None = None,
) -> tuple[Tensor, list[tuple[int, ...]]]:
    r"""Exact forward-filtering backward-sampling of ``num_samples`` derivations from ``p_beta``.

    Returns ``(counts, samples)``: ``counts`` is a ``(num_samples, num_edges)`` tensor of
    per-edge usage counts and ``samples`` the chosen edge-index tuples. Exact (not a
    relaxation); ``counts.mean(0) -> `` :func:`semiring_marginals` as ``num_samples -> inf``.
    Sampling is not differentiable (use :func:`gumbel_relaxed_sample` for a relaxed draw).
    """
    weights = [float(x) for x in edge_weights.detach().reshape(-1).tolist()]
    counts, samples = sample_derivations(graph, weights, beta, num_samples, seed=seed)
    return torch.as_tensor(counts, dtype=edge_weights.dtype, device=edge_weights.device), samples

Related APIsomnibias.struct.sample_derivations↗

Other public import paths

• omnibias.struct.torch.sample_paths
• omnibias.struct.torch.distributions.sample_paths

FUNCTION

topk_free_energy

omnibias.struct.torch.topk_free_energySource ↗

topk_free_energy(graph: Hypergraph, edge_weights: Tensor, k: int, beta: float=1.0) -> Tensor

Differentiable lse_beta restricted to the top-k derivation scores.

The top-k *set* is chosen by exact hard k-best on the detached weights (piecewise constant in the weights); the returned free energy is differentiable through the k selected derivations' scores. -> the best score as beta -> inf and -> the full semiring_value as k -> #derivations (monotone non-decreasing in k).

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · topk free energy 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_topk_free_energy_parity(seed: int) -> None:
    w = _weights(seed)
    wt, wj = torch.tensor(w), jnp.asarray(w)
    for k in (1, 3, 6):
        t = float(st.topk_free_energy(GRAPH, wt, k, 2.0))
        j = float(sj.topk_free_energy(GRAPH, wj, k, 2.0))
        assert abs(t - j) < 1e-10

python -m pytest "packages/omnibias-struct/tests/test_operators.py::test_topk_free_energy_parity[0]" -q

Regression example · topk free energy monotone and limits

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_topk_free_energy_monotone_and_limits(seed: int) -> None:
    w = _weights(seed)
    wt = torch.tensor(w)
    n = count_derivations(GRAPH)
    beta = 2.0
    energies = [float(st.topk_free_energy(GRAPH, wt, k, beta)) for k in range(1, n + 1)]
    assert all(energies[i] <= energies[i + 1] + 1e-12 for i in range(len(energies) - 1))
    assert abs(energies[-1] - float(st.semiring_value(GRAPH, wt, beta))) < 1e-9  # k = N -> full soft value
    best, _ = best_derivation(GRAPH, w)
    assert abs(float(st.topk_free_energy(GRAPH, wt, 1, 128.0)) - best) < 1e-6  # k = 1, beta -> inf

python -m pytest packages/omnibias-struct/tests/test_operators.py -k test_topk_free_energy_monotone_and_limits -q

Read the implementation

def topk_free_energy(graph: Hypergraph, edge_weights: Tensor, k: int, beta: float = 1.0) -> Tensor:
    r"""Differentiable ``lse_beta`` restricted to the top-``k`` derivation scores.

    The top-``k`` *set* is chosen by exact hard k-best on the detached weights (piecewise
    constant in the weights); the returned free energy is differentiable through the ``k``
    selected derivations' scores. ``-> `` the best score as ``beta -> inf`` and ``-> `` the
    full :func:`semiring_value` as ``k -> #derivations`` (monotone non-decreasing in ``k``).
    """
    weights = [float(x) for x in edge_weights.detach().reshape(-1).tolist()]
    kbest = kbest_derivations(graph, weights, k)
    scores = torch.stack(
        [
            torch.stack([edge_weights[e] for e in deriv]).sum() if deriv else edge_weights.new_zeros(())
            for _score, deriv in kbest
        ]
    )
    free_energy: Tensor = logsumexp_beta(scores, beta, axis=-1)
    return free_energy

Related APIsomnibias.struct.kbest_derivations↗omnibias.struct.torch.logsumexp_beta↗

Other public import paths

• omnibias.struct.torch.topk_free_energy
• omnibias.struct.torch.distributions.topk_free_energy

FUNCTION

topk_paths

omnibias.struct.torch.topk_pathsSource ↗

topk_paths(graph: Hypergraph, edge_weights: Tensor, k: int) -> list[tuple[float, tuple[int, ...]]]

Exact k-best derivations (score, edge-tuple) by score, descending (hard decode).

Wraps the exact topological k-best DP; not differentiable (the argmax set is a hard decode). For a smoothed objective over these paths use topk_free_energy.

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnslist[tuple[float, tuple[int, ...]]]

Executed regression · topk paths matches bruteforce

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_topk_paths_matches_bruteforce(seed: int) -> None:
    w = _weights(seed)
    kb = st.topk_paths(GRAPH, torch.tensor(w), 6)
    bk = brute_force_kbest(GRAPH, w, 6)
    assert [round(s, 10) for s, _ in kb] == [round(s, 10) for s, _ in bk]

python -m pytest "packages/omnibias-struct/tests/test_operators.py::test_topk_paths_matches_bruteforce[0]" -q

Read the implementation

def topk_paths(
    graph: Hypergraph, edge_weights: Tensor, k: int
) -> list[tuple[float, tuple[int, ...]]]:
    r"""Exact ``k``-best derivations ``(score, edge-tuple)`` by score, descending (hard decode).

    Wraps the exact topological k-best DP; not differentiable (the argmax set is a hard
    decode). For a smoothed objective over these paths use :func:`topk_free_energy`.
    """
    weights = [float(x) for x in edge_weights.detach().reshape(-1).tolist()]
    result: list[tuple[float, tuple[int, ...]]] = kbest_derivations(graph, weights, k)
    return result

Related APIsomnibias.struct.kbest_derivations↗

Other public import paths

• omnibias.struct.torch.topk_paths
• omnibias.struct.torch.distributions.topk_paths

← Package overviewRunnable learning examples →