API REFERENCE/omnibias-struct

omnibias.struct.torch.soft_dp

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.torch.soft_dp

Module contract & mathematical conventions

Differentiable soft dynamic-programming layers (torch).

Algebraically corresponding twin of omnibias.struct.jax.soft_dp (float64). Each soft_* replaces the hard max / min combine of its DP recursion with lse_beta (omnibias.struct.torch._logsumexp), unrolled so autograd flows through it. The beta -> inf annealing is the *feasibility / temperature* collapse (a soft object hardened to a discrete one, the same axis as omnibias-discrete / -qubo / -routing); it is differentiated exactly by the delta -> 0 founding bias collapse tower -- do not conflate the two axes. The *_marginals functions return the **closed-form** gradient -- the forward-backward path marginal assembled from the tower's softmax -- which the tests pin equal to autograd.

**Numerical envelope (measured).** Use float64: the soft *value* stays bit-stable out to beta ~ 1e12 and the closed-form marginals are accurate to ~1e-9 for beta up to about 1e6, degrading gracefully (never NaN) beyond. The unreachable-cell sentinel _NEG is safe in float64 (beta * _NEG never overflows for practical beta); float32 breaks both algebraically correspondingity and large-beta stability. The *_batched helpers map a leading batch dimension through torch.func.vmap and are algebraically corresponding to looping the per-example layer.

Numerical scope: shared coefficients and formulas do not establish universal floating-point identity across backends, devices, compiler transforms, or reduction orders. Compare the chosen path with dtype-appropriate tolerances; even float64 is not a general bitwise-equality guarantee. Algebraic initialization identities likewise need a numerical check for the selected coefficients and reduction order.

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

soft_ctc

omnibias.struct.torch.soft_ctcSource ↗

soft_ctc(log_probs: Tensor, lattice: CTCLattice, beta: float=1.0) -> Tensor

Soft CTC value beta^-1 log sum_alignments exp(beta score) for log_probs (T, C).

At beta = 1 this is the standard CTC log-likelihood; as beta -> inf it anneals to the best single-alignment score. Differentiable in log_probs (autograd through the unrolled lse_beta lattice recursion; its gradient is the alignment marginal).

Parameters

Name	Type & default	Contract
log_probspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
latticepositional-or-keyword	CTCLatticeRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · closed form marginals match autodiff

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_closed_form_marginals_match_autodiff(problem: str, backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_marginals_vs_autodiff(problem, 0, 6.0, backend)
    assert rec.ok, rec

python -m pytest "packages/omnibias-struct/tests/test_refinement_probes.py::test_closed_form_marginals_match_autodiff[torch-ctc]" -q

Regression example · soft ctc 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_ctc_parity(seed: int) -> None:
    lattice, log_probs = sample_ctc(seed)
    lpt, lpj = torch.tensor(log_probs), jnp.asarray(log_probs)
    for beta in (1.0, 8.0):
        assert abs(float(st.soft_ctc(lpt, lattice, beta)) - float(sj.soft_ctc(lpj, lattice, beta))) < ACROSS_SEED_TOL

python -m pytest packages/omnibias-struct/tests/test_parity.py -k test_soft_ctc_parity -q

Read the implementation

def soft_ctc(log_probs: Tensor, lattice: CTCLattice, beta: float = 1.0) -> Tensor:
    r"""Soft CTC value ``beta^-1 log sum_alignments exp(beta score)`` for ``log_probs`` ``(T, C)``.

    At ``beta = 1`` this is the standard CTC log-likelihood; as ``beta -> inf`` it anneals
    to the best single-alignment score. Differentiable in ``log_probs`` (autograd through
    the unrolled ``lse_beta`` lattice recursion; its gradient is the alignment marginal).
    """
    ext = lattice.extended_labels()
    m = int(ext.shape[0])
    n_steps = log_probs.shape[0]
    neg = torch.full((), _NEG, dtype=log_probs.dtype, device=log_probs.device)
    f: list[Tensor] = []
    for s in range(m):
        if s == 0:
            f.append(log_probs[0, int(ext[0])])
        elif s == 1:
            f.append(log_probs[0, int(ext[1])])
        else:
            f.append(neg)
    for t in range(1, n_steps):
        nf: list[Tensor] = []
        for s in range(m):
            vals = torch.stack([f[p] for p in lattice.incoming(s)])
            nf.append(log_probs[t, int(ext[s])] + logsumexp_beta(vals, beta, axis=-1))
        f = nf
    ends = [f[m - 1]] + ([f[m - 2]] if m >= 2 else [])
    return logsumexp_beta(torch.stack(ends), beta, axis=-1)

Related APIsomnibias.struct.torch.logsumexp_beta↗

Other public import paths

• omnibias.struct.torch.soft_ctc
• omnibias.struct.torch.soft_dp.soft_ctc

FUNCTION

soft_ctc_batched

omnibias.struct.torch.soft_ctc_batchedSource ↗

soft_ctc_batched(log_probs: Tensor, lattice: CTCLattice, beta: float=1.0) -> Tensor

Batched soft_ctc -> (B,) for log_probs (B, T, C) (shared lattice).

Parameters

Name	Type & default	Contract
log_probspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
latticepositional-or-keyword	CTCLatticeRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft shortest path and ctc batched torch

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_shortest_path_and_ctc_batched_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import (
        soft_ctc,
        soft_ctc_batched,
        soft_shortest_path,
        soft_shortest_path_batched,
    )

    dag = random_dag(0)
    weights = torch.stack([torch.tensor(dag_weight_matrix(dag)) * (1.0 + 0.1 * i) for i in range(4)])
    b = soft_shortest_path_batched(weights, dag, 6.0)
    loop = torch.stack([soft_shortest_path(weights[i], dag, 6.0) for i in range(4)])
    assert torch.max(torch.abs(b - loop)).item() < ATOL

    lattice, lp = sample_ctc(0)
    log_probs = torch.stack([torch.tensor(lp) + 0.05 * i for i in range(4)])
    bc = soft_ctc_batched(log_probs, lattice, 4.0)
    loopc = torch.stack([soft_ctc(log_probs[i], lattice, 4.0) for i in range(4)])
    assert torch.max(torch.abs(bc - loopc)).item() < ATOL

python -m pytest "packages/omnibias-struct/tests/test_batch.py::test_soft_shortest_path_and_ctc_batched_torch" -q

Read the implementation

def soft_ctc_batched(log_probs: Tensor, lattice: CTCLattice, beta: float = 1.0) -> Tensor:
    r"""Batched :func:`soft_ctc` -> ``(B,)`` for ``log_probs`` ``(B, T, C)`` (shared ``lattice``)."""
    from torch.func import vmap

    def fwd(lp: Tensor) -> Tensor:
        return soft_ctc(lp, lattice, beta)

    out: Tensor = vmap(fwd)(log_probs)
    return out

Related APIsomnibias.struct.torch.soft_ctc↗

Other public import paths

• omnibias.struct.torch.soft_ctc_batched
• omnibias.struct.torch.soft_dp.soft_ctc_batched

FUNCTION

soft_ctc_marginals

omnibias.struct.torch.soft_ctc_marginalsSource ↗

soft_ctc_marginals(log_probs: Tensor, lattice: CTCLattice, beta: float=1.0) -> Tensor

Closed-form CTC label marginals d soft_ctc / d log_probs ((T, C)).

Forward-backward over the blank-augmented lattice with the tower softmax; entry [t, c] is the total probability mass of alignments emitting class c at time t, so each row sums to 1. Equals autograd of soft_ctc and, as beta -> inf, concentrates on the best alignment (ctc_best_alignment).

Parameters

Name	Type & default	Contract
log_probspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
latticepositional-or-keyword	CTCLatticeRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · closed form marginals match autodiff

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_closed_form_marginals_match_autodiff(problem: str, backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_marginals_vs_autodiff(problem, 0, 6.0, backend)
    assert rec.ok, rec

python -m pytest "packages/omnibias-struct/tests/test_refinement_probes.py::test_closed_form_marginals_match_autodiff[torch-ctc]" -q

Regression example · ctc marginals concentrate on best alignment

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_ctc_marginals_concentrate_on_best_alignment() -> None:
    pytest.importorskip("torch")
    import torch
    from _struct_helpers import sample_ctc
    from omnibias.struct import ctc_best_alignment
    from omnibias.struct.torch import soft_ctc_marginals

    lattice, log_probs = sample_ctc(3)
    _, alignment = ctc_best_alignment(lattice, log_probs)
    marg = soft_ctc_marginals(torch.tensor(log_probs), lattice, 256.0).numpy()
    assert tuple(int(c) for c in marg.argmax(axis=1)) == alignment

python -m pytest packages/omnibias-struct/tests/test_gradients.py -k test_ctc_marginals_concentrate_on_best_alignment -q

Read the implementation

def soft_ctc_marginals(log_probs: Tensor, lattice: CTCLattice, beta: float = 1.0) -> Tensor:
    r"""Closed-form CTC label marginals ``d soft_ctc / d log_probs`` (``(T, C)``).

    Forward-backward over the blank-augmented lattice with the tower softmax; entry
    ``[t, c]`` is the total probability mass of alignments emitting class ``c`` at time
    ``t``, so each row sums to ``1``. Equals ``autograd`` of :func:`soft_ctc` and, as
    ``beta -> inf``, concentrates on the best alignment (:func:`ctc_best_alignment`).
    """
    lp = log_probs
    ext = lattice.extended_labels()
    m = int(ext.shape[0])
    n_steps = lp.shape[0]
    n_classes = int(lp.shape[1])
    neg = torch.full((), _NEG, dtype=lp.dtype, device=lp.device)
    zero = torch.zeros((), dtype=lp.dtype, device=lp.device)
    # forward alpha (includes the current-cell emission)
    f: list[Tensor] = [lp[0, int(ext[s])] if s < 2 else neg for s in range(m)]
    alpha_rows = [torch.stack(f)]
    for t in range(1, n_steps):
        nf = [
            lp[t, int(ext[s])] + logsumexp_beta(torch.stack([f[p] for p in lattice.incoming(s)]), beta, axis=-1)
            for s in range(m)
        ]
        f = nf
        alpha_rows.append(torch.stack(f))
    alpha = torch.stack(alpha_rows, dim=0)  # (T, m)
    # outgoing edges of the lattice (transpose of incoming)
    outgoing: dict[int, list[int]] = {s: [] for s in range(m)}
    for s in range(m):
        for p in lattice.incoming(s):
            outgoing[p].append(s)
    # backward beta (suffix score, excludes the current-cell emission)
    b: list[Tensor] = [zero if s in (m - 1, m - 2) else neg for s in range(m)]
    bwd_rows = [torch.stack(b)]
    for t in range(n_steps - 2, -1, -1):
        nb = [
            logsumexp_beta(
                torch.stack([lp[t + 1, int(ext[s2])] + b[s2] for s2 in outgoing[s]]), beta, axis=-1
            )
            if outgoing[s]
            else neg
            for s in range(m)
        ]
        b = nb
        bwd_rows.insert(0, torch.stack(b))
    bwd = torch.stack(bwd_rows, dim=0)  # (T, m)
    end_vals = [alpha_rows[n_steps - 1][m - 1]] + ([alpha_rows[n_steps - 1][m - 2]] if m >= 2 else [])
    value = logsumexp_beta(torch.stack(end_vals), beta, axis=-1)
    gamma_cell = torch.exp(beta * (alpha + bwd - value))  # (T, m)
    cols: list[Tensor] = []
    for c in range(n_classes):
        members = [s for s in range(m) if int(ext[s]) == c]
        col = torch.zeros(n_steps, dtype=lp.dtype, device=lp.device)
        for s in members:
            col = col + gamma_cell[:, s]
        cols.append(col)
    return torch.stack(cols, dim=1)  # (T, C)

Related APIsomnibias.struct.torch.logsumexp_beta↗

Other public import paths

• omnibias.struct.torch.soft_ctc_marginals
• omnibias.struct.torch.soft_dp.soft_ctc_marginals

FUNCTION

soft_shortest_path

omnibias.struct.torch.soft_shortest_pathSource ↗

soft_shortest_path(weights: Tensor, dag: DAG, beta: float=1.0) -> Tensor

Soft (softmin) shortest-path cost of a DAG for edge weights (n, n).

Returns -beta^-1 log sum_paths exp(-beta cost) (softmin over source -> sink path costs); -> min cost as beta -> inf. Differentiable in the finite edge entries of weights. Certify in the max convention via the negated value (V* = -min cost, V_beta = -soft cost).

Parameters

Name	Type & default	Contract
weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
dagpositional-or-keyword	DAGRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When sink_value is None, raises ValueError('sink is not reachable from source in this DAG').

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

Executed regression · alignment soft ge hard

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_alignment_soft_ge_hard(kind: str) -> None:
    pytest.importorskip("torch")
    assert probe_align_gap(kind, 0, BETA, "torch").ok

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

Regression example · soft shortest path 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_shortest_path_parity(seed: int) -> None:
    dag = random_dag(seed)
    wt, wj = torch.tensor(dag_weight_matrix(dag)), jnp.asarray(dag_weight_matrix(dag))
    for beta in (1.0, 8.0):
        assert abs(float(st.soft_shortest_path(wt, dag, beta)) - float(sj.soft_shortest_path(wj, dag, beta))) < ACROSS_SEED_TOL
        xt = st.soft_shortest_path_marginals(wt, dag, beta)
        xj = sj.soft_shortest_path_marginals(wj, dag, beta)
        assert _max_abs_diff(xt, xj) < ACROSS_SEED_TOL

python -m pytest packages/omnibias-struct/tests/test_parity.py -k test_soft_shortest_path_parity -q

Read the implementation

def soft_shortest_path(weights: Tensor, dag: DAG, beta: float = 1.0) -> Tensor:
    r"""Soft (softmin) shortest-path cost of a :class:`DAG` for edge ``weights`` ``(n, n)``.

    Returns ``-beta^-1 log sum_paths exp(-beta cost)`` (softmin over ``source -> sink``
    path costs); ``-> min cost`` as ``beta -> inf``. Differentiable in the finite edge
    entries of ``weights``. Certify in the ``max`` convention via the negated value
    (``V* = -min cost``, ``V_beta = -soft cost``).
    """
    alpha = _dag_forward(weights, dag, beta)
    sink_value = alpha[dag.sink]
    if sink_value is None:
        raise ValueError("sink is not reachable from source in this DAG")
    return -sink_value

Other public import paths

• omnibias.struct.torch.soft_shortest_path
• omnibias.struct.torch.soft_dp.soft_shortest_path

FUNCTION

soft_shortest_path_batched

omnibias.struct.torch.soft_shortest_path_batchedSource ↗

soft_shortest_path_batched(weights: Tensor, dag: DAG, beta: float=1.0) -> Tensor

Batched soft_shortest_path -> (B,) for weights (B, n, n) (shared dag).

Parameters

Name	Type & default	Contract
weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
dagpositional-or-keyword	DAGRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft shortest path and ctc batched torch

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_shortest_path_and_ctc_batched_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import (
        soft_ctc,
        soft_ctc_batched,
        soft_shortest_path,
        soft_shortest_path_batched,
    )

    dag = random_dag(0)
    weights = torch.stack([torch.tensor(dag_weight_matrix(dag)) * (1.0 + 0.1 * i) for i in range(4)])
    b = soft_shortest_path_batched(weights, dag, 6.0)
    loop = torch.stack([soft_shortest_path(weights[i], dag, 6.0) for i in range(4)])
    assert torch.max(torch.abs(b - loop)).item() < ATOL

    lattice, lp = sample_ctc(0)
    log_probs = torch.stack([torch.tensor(lp) + 0.05 * i for i in range(4)])
    bc = soft_ctc_batched(log_probs, lattice, 4.0)
    loopc = torch.stack([soft_ctc(log_probs[i], lattice, 4.0) for i in range(4)])
    assert torch.max(torch.abs(bc - loopc)).item() < ATOL

python -m pytest "packages/omnibias-struct/tests/test_batch.py::test_soft_shortest_path_and_ctc_batched_torch" -q

Read the implementation

def soft_shortest_path_batched(weights: Tensor, dag: DAG, beta: float = 1.0) -> Tensor:
    r"""Batched :func:`soft_shortest_path` -> ``(B,)`` for ``weights`` ``(B, n, n)`` (shared ``dag``)."""
    from torch.func import vmap

    def fwd(w: Tensor) -> Tensor:
        return soft_shortest_path(w, dag, beta)

    out: Tensor = vmap(fwd)(weights)
    return out

Related APIsomnibias.struct.torch.soft_shortest_path↗

Other public import paths

• omnibias.struct.torch.soft_shortest_path_batched
• omnibias.struct.torch.soft_dp.soft_shortest_path_batched

FUNCTION

soft_shortest_path_marginals

omnibias.struct.torch.soft_shortest_path_marginalsSource ↗

soft_shortest_path_marginals(weights: Tensor, dag: DAG, beta: float=1.0) -> Tensor

Closed-form edge-usage marginals xi[u, v] = P_beta(edge (u, v) on the path).

An (n, n) tensor, zero off the edge set. Equals d soft_shortest_path / d weights on the edges; forward-backward with the tower softmax. As beta -> inf it concentrates on the shortest path.

Parameters

Name	Type & default	Contract
weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
dagpositional-or-keyword	DAGRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When value is None, raises ValueError('sink is not reachable from source in this DAG').

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

Executed regression · closed form marginals match autodiff

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_closed_form_marginals_match_autodiff(problem: str, backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_marginals_vs_autodiff(problem, 0, 6.0, backend)
    assert rec.ok, rec

python -m pytest "packages/omnibias-struct/tests/test_refinement_probes.py::test_closed_form_marginals_match_autodiff[torch-shortest_path]" -q

Regression example · soft shortest path 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_shortest_path_parity(seed: int) -> None:
    dag = random_dag(seed)
    wt, wj = torch.tensor(dag_weight_matrix(dag)), jnp.asarray(dag_weight_matrix(dag))
    for beta in (1.0, 8.0):
        assert abs(float(st.soft_shortest_path(wt, dag, beta)) - float(sj.soft_shortest_path(wj, dag, beta))) < ACROSS_SEED_TOL
        xt = st.soft_shortest_path_marginals(wt, dag, beta)
        xj = sj.soft_shortest_path_marginals(wj, dag, beta)
        assert _max_abs_diff(xt, xj) < ACROSS_SEED_TOL

python -m pytest packages/omnibias-struct/tests/test_parity.py -k test_soft_shortest_path_parity -q

Read the implementation

def soft_shortest_path_marginals(weights: Tensor, dag: DAG, beta: float = 1.0) -> Tensor:
    r"""Closed-form edge-usage marginals ``xi[u, v] = P_beta(edge (u, v) on the path)``.

    An ``(n, n)`` tensor, zero off the edge set. Equals ``d soft_shortest_path / d
    weights`` on the edges; forward-backward with the tower softmax. As ``beta -> inf``
    it concentrates on the shortest path.
    """
    alpha = _dag_forward(weights, dag, beta)
    bwd = _dag_backward(weights, dag, beta)
    value = alpha[dag.sink]
    if value is None:
        raise ValueError("sink is not reachable from source in this DAG")
    xi = torch.zeros_like(weights)
    for u, v in dag.edges:
        au, bv = alpha[u], bwd[v]
        if au is None or bv is None:
            continue
        xi[u, v] = torch.exp(beta * (au - weights[u, v] + bv - value))
    return xi

Other public import paths

• omnibias.struct.torch.soft_shortest_path_marginals
• omnibias.struct.torch.soft_dp.soft_shortest_path_marginals

FUNCTION

soft_viterbi

omnibias.struct.torch.soft_viterbiSource ↗

soft_viterbi(emissions: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Soft Viterbi value beta^-1 log sum_paths exp(beta score) of a linear chain.

emissions is (T, S), transitions is (S, S), optional start is (S,). Differentiable in all three; -> max_path score as beta -> inf.

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · closed form marginals match autodiff

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_closed_form_marginals_match_autodiff(problem: str, backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_marginals_vs_autodiff(problem, 0, 6.0, backend)
    assert rec.ok, rec

python -m pytest "packages/omnibias-struct/tests/test_refinement_probes.py::test_closed_form_marginals_match_autodiff[torch-viterbi]" -q

Regression example · soft viterbi equals log partition

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_viterbi_equals_log_partition(backend: tuple[Any, Callable[[Any], Any]]) -> None:
    mod, to_t = backend
    for seed in range(4):
        trellis = random_chain(seed)
        emissions = to_t(trellis.emissions)
        transitions = to_t(trellis.transitions)
        start = to_t(trellis.start)
        for beta in (1.0, 4.0, 16.0):
            soft = float(mod.soft_viterbi(emissions, transitions, beta, start=start))
            assert abs(soft - brute_force_partition(trellis, beta)) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_soft_dp.py -k test_soft_viterbi_equals_log_partition -q

Read the implementation

def soft_viterbi(
    emissions: Tensor,
    transitions: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Soft Viterbi value ``beta^-1 log sum_paths exp(beta score)`` of a linear chain.

    ``emissions`` is ``(T, S)``, ``transitions`` is ``(S, S)``, optional ``start`` is
    ``(S,)``. Differentiable in all three; ``-> max_path score`` as ``beta -> inf``.
    """
    start_t = _resolve_start(start, emissions.shape[1], emissions)
    alpha = _chain_forward(emissions, transitions, start_t, beta)
    return logsumexp_beta(alpha[-1], beta, axis=-1)

Related APIsomnibias.struct.torch.logsumexp_beta↗

Other public import paths

• omnibias.struct.torch.soft_viterbi
• omnibias.struct.torch.soft_dp.soft_viterbi

FUNCTION

soft_viterbi_batched

omnibias.struct.torch.soft_viterbi_batchedSource ↗

soft_viterbi_batched(emissions: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Batched soft_viterbi -> (B,) for emissions (B, T, S).

transitions may be shared (S, S) or per-example (B, S, S); start may be None, shared (S,), or per-example (B, S). Maps soft_viterbi over the leading batch axis with torch.func.vmap -- bit-identical to looping.

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft viterbi batched matches loop and shapes torch

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_viterbi_batched_matches_loop_and_shapes_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi, soft_viterbi_batched

    trellises = [random_chain(s) for s in range(4)]
    emissions = torch.stack([torch.tensor(t.emissions) for t in trellises])
    transitions = torch.tensor(trellises[0].transitions)  # shared
    start = torch.tensor(trellises[0].start)
    batched = soft_viterbi_batched(emissions, transitions, 5.0, start=start)
    loop = torch.stack([soft_viterbi(emissions[i], transitions, 5.0, start=start) for i in range(4)])
    assert batched.shape == (4,)
    assert torch.max(torch.abs(batched - loop)).item() < ATOL

python -m pytest "packages/omnibias-struct/tests/test_batch.py::test_soft_viterbi_batched_matches_loop_and_shapes_torch" -q

Regression example · soft viterbi batched per example transitions and start torch

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_viterbi_batched_per_example_transitions_and_start_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi, soft_viterbi_batched

    trellises = [random_chain(s) for s in range(3)]
    emissions = torch.stack([torch.tensor(t.emissions) for t in trellises])
    transitions = torch.stack([torch.tensor(t.transitions) for t in trellises])
    start = torch.stack([torch.tensor(t.start) for t in trellises])
    batched = soft_viterbi_batched(emissions, transitions, 3.0, start=start)
    loop = torch.stack(
        [soft_viterbi(emissions[i], transitions[i], 3.0, start=start[i]) for i in range(3)]
    )
    assert torch.max(torch.abs(batched - loop)).item() < ATOL

python -m pytest packages/omnibias-struct/tests/test_batch.py -k test_soft_viterbi_batched_per_example_transitions_and_start_torch -q

Read the implementation

def soft_viterbi_batched(
    emissions: Tensor,
    transitions: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Batched :func:`soft_viterbi` -> ``(B,)`` for ``emissions`` ``(B, T, S)``.

    ``transitions`` may be shared ``(S, S)`` or per-example ``(B, S, S)``; ``start`` may be
    ``None``, shared ``(S,)``, or per-example ``(B, S)``. Maps :func:`soft_viterbi` over
    the leading batch axis with ``torch.func.vmap`` -- bit-identical to looping.
    """
    from torch.func import vmap

    t_dim = 0 if transitions.dim() == 3 else None
    if start is None:
        def fwd(e: Tensor, tr: Tensor) -> Tensor:
            return soft_viterbi(e, tr, beta)

        out: Tensor = vmap(fwd, in_dims=(0, t_dim))(emissions, transitions)
        return out
    s_dim = 0 if start.dim() == 2 else None

    def fwd_start(e: Tensor, tr: Tensor, st: Tensor) -> Tensor:
        return soft_viterbi(e, tr, beta, start=st)

    out_s: Tensor = vmap(fwd_start, in_dims=(0, t_dim, s_dim))(emissions, transitions, start)
    return out_s

Related APIsomnibias.struct.torch.soft_viterbi↗

Other public import paths

• omnibias.struct.torch.soft_viterbi_batched
• omnibias.struct.torch.soft_dp.soft_viterbi_batched

FUNCTION

soft_viterbi_marginals

omnibias.struct.torch.soft_viterbi_marginalsSource ↗

soft_viterbi_marginals(emissions: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Closed-form state marginals gamma[t, s] = P_beta(state s at time t) ((T, S)).

Forward-backward with the tower softmax; equals d soft_viterbi / d emissions and sums to 1 along the state axis. As beta -> inf it concentrates on the Viterbi path.

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · closed form marginals match autodiff

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_closed_form_marginals_match_autodiff(problem: str, backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_marginals_vs_autodiff(problem, 0, 6.0, backend)
    assert rec.ok, rec

python -m pytest "packages/omnibias-struct/tests/test_refinement_probes.py::test_closed_form_marginals_match_autodiff[torch-viterbi]" -q

Regression example · viterbi marginals accurate through beta 1e6

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_viterbi_marginals_accurate_through_beta_1e6() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi_marginals

    trellis = random_chain(0)
    for beta in (1.0e2, 1.0e4, 1.0e6):
        gamma = soft_viterbi_marginals(
            torch.tensor(trellis.emissions), torch.tensor(trellis.transitions), beta,
            start=torch.tensor(trellis.start),
        ).numpy()
        assert np.max(np.abs(gamma.sum(axis=1) - 1.0)) < 1e-8

python -m pytest packages/omnibias-struct/tests/test_batch.py -k test_viterbi_marginals_accurate_through_beta_1e6 -q

Read the implementation

def soft_viterbi_marginals(
    emissions: Tensor,
    transitions: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Closed-form state marginals ``gamma[t, s] = P_beta(state s at time t)`` (``(T, S)``).

    Forward-backward with the tower softmax; equals ``d soft_viterbi / d emissions`` and
    sums to ``1`` along the state axis. As ``beta -> inf`` it concentrates on the Viterbi
    path.
    """
    start_t = _resolve_start(start, emissions.shape[1], emissions)
    alpha = _chain_forward(emissions, transitions, start_t, beta)  # (T, S)
    bwd = _chain_backward(emissions, transitions, beta)  # (T, S)
    value = logsumexp_beta(alpha[-1], beta, axis=-1)
    return torch.exp(beta * (alpha + bwd - value))

Related APIsomnibias.struct.torch.logsumexp_beta↗

Other public import paths

• omnibias.struct.torch.soft_viterbi_marginals
• omnibias.struct.torch.soft_dp.soft_viterbi_marginals

FUNCTION

soft_viterbi_transition_marginals

omnibias.struct.torch.soft_viterbi_transition_marginalsSource ↗

soft_viterbi_transition_marginals(emissions: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Closed-form transition-pair marginals xi[t, s, s'] ((T - 1, S, S)).

xi[t, s, s'] = P_beta(state s at time t, state s' at time t + 1). Each (S, S) slab sums to 1, and summing over t gives d soft_viterbi / d transitions -- the transition analogue of soft_viterbi_marginals. As beta -> inf it concentrates on the consecutive-state pairs of the Viterbi path.

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · transition marginals equal autograd torch

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_transition_marginals_equal_autograd_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi, soft_viterbi_transition_marginals

    trellis = random_chain(1)
    emissions = torch.tensor(trellis.emissions)
    transitions = torch.tensor(trellis.transitions, requires_grad=True)
    start = torch.tensor(trellis.start)
    soft_viterbi(emissions, transitions, 6.0, start=start).backward()
    xi = soft_viterbi_transition_marginals(emissions, transitions.detach(), 6.0, start=start)
    assert np.allclose(xi.sum(dim=0).numpy(), transitions.grad.numpy(), atol=1e-9)
    assert np.allclose(xi.numpy().sum(axis=(1, 2)), 1.0, atol=1e-9)  # each (S,S) slab sums to 1

python -m pytest "packages/omnibias-struct/tests/test_batch.py::test_transition_marginals_equal_autograd_torch" -q

Regression example · transition marginals torch jax 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_transition_marginals_torch_jax_parity() -> None:
    pytest.importorskip("torch")
    pytest.importorskip("jax")
    import jax
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    import torch
    from omnibias.struct.jax import soft_viterbi_transition_marginals as jax_xi
    from omnibias.struct.torch import soft_viterbi_transition_marginals as torch_xi

    trellis = random_chain(2)
    xt = torch_xi(
        torch.tensor(trellis.emissions), torch.tensor(trellis.transitions), 5.0,
        start=torch.tensor(trellis.start),
    ).numpy()
    xj = np.asarray(
        jax_xi(
            jnp.asarray(trellis.emissions), jnp.asarray(trellis.transitions), 5.0,
            start=jnp.asarray(trellis.start),
        )
    )
    assert np.max(np.abs(xt - xj)) < ATOL

python -m pytest packages/omnibias-struct/tests/test_batch.py -k test_transition_marginals_torch_jax_parity -q

Read the implementation

def soft_viterbi_transition_marginals(
    emissions: Tensor,
    transitions: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Closed-form transition-pair marginals ``xi[t, s, s']`` (``(T - 1, S, S)``).

    ``xi[t, s, s'] = P_beta(state s at time t, state s' at time t + 1)``. Each ``(S, S)``
    slab sums to ``1``, and summing over ``t`` gives ``d soft_viterbi / d transitions`` --
    the transition analogue of :func:`soft_viterbi_marginals`. As ``beta -> inf`` it
    concentrates on the consecutive-state pairs of the Viterbi path.
    """
    start_t = _resolve_start(start, emissions.shape[1], emissions)
    n_steps = emissions.shape[0]
    alpha = _chain_forward(emissions, transitions, start_t, beta)  # (T, S)
    bwd = _chain_backward(emissions, transitions, beta)  # (T, S)
    value = logsumexp_beta(alpha[-1], beta, axis=-1)
    xis: list[Tensor] = []
    for t in range(1, n_steps):
        # alpha[t-1][s] + A[s, s'] + (E[t][s'] + B[t][s'])
        term = alpha[t - 1][:, None] + transitions + (emissions[t] + bwd[t])[None, :]
        xis.append(torch.exp(beta * (term - value)))
    return torch.stack(xis, dim=0)  # (T - 1, S, S)

Related APIsomnibias.struct.torch.logsumexp_beta↗

Other public import paths

• omnibias.struct.torch.soft_viterbi_transition_marginals
• omnibias.struct.torch.soft_dp.soft_viterbi_transition_marginals

← Package overviewRunnable learning examples →