API REFERENCE/omnibias-struct

omnibias.struct.jax.soft_dp

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.soft_dp

Module contract & mathematical conventions

Differentiable soft dynamic-programming layers (jax).

Algebraically corresponding twin of omnibias.struct.torch.soft_dp (float64 -- enable jax_enable_x64). Each soft_* replaces the hard max / min combine of its DP recursion with lse_beta (omnibias.struct.jax._logsumexp), unrolled so jax.grad 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 jax.grad. Control flow branches only on the static problem structure, so the functions stay grad-safe.

**Numerical envelope (measured).** Use float64 (jax_enable_x64): 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; float32 breaks both algebraically correspondingity and large-beta stability. The *_batched helpers map a leading batch dimension through jax.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.jax.soft_ctcSource ↗

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

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 (grad through the unrolled lse_beta lattice recursion; its gradient is the alignment marginal).

Parameters

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

ReturnsArray

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[jax-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: Array, lattice: CTCLattice, beta: float = 1.0) -> Array:
    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`` (grad 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 = jnp.full((), _NEG, dtype=log_probs.dtype)
    f: list[Array] = []
    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[Array] = []
        for s in range(m):
            vals = jnp.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(jnp.stack(ends), beta, axis=-1)

Related APIsomnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.soft_ctc
• omnibias.struct.jax.soft_dp.soft_ctc

FUNCTION

soft_ctc_batched

omnibias.struct.jax.soft_ctc_batchedSource ↗

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

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

Parameters

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

ReturnsArray

Read the implementation

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

    out: Array = jax.vmap(lambda lp: soft_ctc(lp, lattice, beta))(log_probs)
    return out

Related APIsomnibias.struct.jax.soft_ctc↗

Other public import paths

• omnibias.struct.jax.soft_ctc_batched
• omnibias.struct.jax.soft_dp.soft_ctc_batched

FUNCTION

soft_ctc_marginals

omnibias.struct.jax.soft_ctc_marginalsSource ↗

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

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 jax.grad of soft_ctc and, as beta -> inf, concentrates on the best alignment.

Parameters

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

ReturnsArray

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[jax-ctc]" -q

Regression example · ctc 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_ctc_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 _struct_helpers import sample_ctc
    from omnibias.struct.jax import soft_ctc_marginals as jax_marg
    from omnibias.struct.torch import soft_ctc_marginals as torch_marg

    lattice, log_probs = sample_ctc(2)
    mt = torch_marg(torch.tensor(log_probs), lattice, 5.0).numpy()
    mj = np.asarray(jax_marg(jnp.asarray(log_probs), lattice, 5.0))
    assert np.max(np.abs(mt - mj)) < 1e-12

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

Read the implementation

def soft_ctc_marginals(log_probs: Array, lattice: CTCLattice, beta: float = 1.0) -> Array:
    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 ``jax.grad`` of :func:`soft_ctc` and, as
    ``beta -> inf``, concentrates on the 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 = jnp.full((), _NEG, dtype=lp.dtype)
    zero = jnp.zeros((), dtype=lp.dtype)
    # forward alpha (includes the current-cell emission)
    f: list[Array] = [lp[0, int(ext[s])] if s < 2 else neg for s in range(m)]
    alpha_rows = [jnp.stack(f)]
    for t in range(1, n_steps):
        nf = [
            lp[t, int(ext[s])] + logsumexp_beta(jnp.stack([f[p] for p in lattice.incoming(s)]), beta, axis=-1)
            for s in range(m)
        ]
        f = nf
        alpha_rows.append(jnp.stack(f))
    alpha = jnp.stack(alpha_rows, axis=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[Array] = [zero if s in (m - 1, m - 2) else neg for s in range(m)]
    bwd_rows = [jnp.stack(b)]
    for t in range(n_steps - 2, -1, -1):
        nb = [
            logsumexp_beta(
                jnp.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, jnp.stack(b))
    bwd = jnp.stack(bwd_rows, axis=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(jnp.stack(end_vals), beta, axis=-1)
    gamma_cell = jnp.exp(beta * (alpha + bwd - value))  # (T, m)
    cols: list[Array] = []
    for c in range(n_classes):
        members = [s for s in range(m) if int(ext[s]) == c]
        col = jnp.zeros(n_steps, dtype=lp.dtype)
        for s in members:
            col = col + gamma_cell[:, s]
        cols.append(col)
    return jnp.stack(cols, axis=1)  # (T, C)

Related APIsomnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.soft_ctc_marginals
• omnibias.struct.jax.soft_dp.soft_ctc_marginals

FUNCTION

soft_shortest_path

omnibias.struct.jax.soft_shortest_pathSource ↗

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

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	ArrayRequired	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.

ReturnsArray

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 parity bit identical

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_parity_bit_identical(kind: str) -> None:
    pytest.importorskip("torch")
    pytest.importorskip("jax")
    assert probe_align_parity(kind, 0, BETA).metrics["value_abs_diff"] < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_alignment_parity_bit_identical[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: Array, dag: DAG, beta: float = 1.0) -> Array:
    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.jax.soft_shortest_path
• omnibias.struct.jax.soft_dp.soft_shortest_path

FUNCTION

soft_shortest_path_batched

omnibias.struct.jax.soft_shortest_path_batchedSource ↗

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

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

Parameters

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

ReturnsArray

Read the implementation

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

    out: Array = jax.vmap(lambda w: soft_shortest_path(w, dag, beta))(weights)
    return out

Related APIsomnibias.struct.jax.soft_shortest_path↗

Other public import paths

• omnibias.struct.jax.soft_shortest_path_batched
• omnibias.struct.jax.soft_dp.soft_shortest_path_batched

FUNCTION

soft_shortest_path_marginals

omnibias.struct.jax.soft_shortest_path_marginalsSource ↗

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

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

An (n, n) array, 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	ArrayRequired	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.

ReturnsArray

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[jax-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: Array, dag: DAG, beta: float = 1.0) -> Array:
    r"""Closed-form edge-usage marginals ``xi[u, v] = P_beta(edge (u, v) on the path)``.

    An ``(n, n)`` array, 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 = jnp.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 = xi.at[u, v].set(jnp.exp(beta * (au - weights[u, v] + bv - value)))
    return xi

Other public import paths

• omnibias.struct.jax.soft_shortest_path_marginals
• omnibias.struct.jax.soft_dp.soft_shortest_path_marginals

FUNCTION

soft_viterbi

omnibias.struct.jax.soft_viterbiSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

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[jax-viterbi]" -q

Regression example · viterbi marginals equal autograd 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_viterbi_marginals_equal_autograd_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi, soft_viterbi_marginals

    trellis = random_chain(0)
    emissions = torch.tensor(trellis.emissions, requires_grad=True)
    transitions = torch.tensor(trellis.transitions)
    start = torch.tensor(trellis.start)
    soft_viterbi(emissions, transitions, 6.0, start=start).backward()
    gamma = soft_viterbi_marginals(emissions.detach(), transitions, 6.0, start=start)
    assert np.allclose(emissions.grad.numpy(), gamma.numpy(), atol=ATOL)
    assert np.allclose(gamma.numpy().sum(axis=1), 1.0, atol=ATOL)

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

Read the implementation

def soft_viterbi(
    emissions: Array,
    transitions: Array,
    beta: float = 1.0,
    *,
    start: Array | None = None,
) -> Array:
    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.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.soft_viterbi
• omnibias.struct.jax.soft_dp.soft_viterbi

FUNCTION

soft_viterbi_batched

omnibias.struct.jax.soft_viterbi_batchedSource ↗

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

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 jax.vmap -- bit-identical to looping.

Parameters

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

ReturnsArray

Executed regression · batched torch jax 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_batched_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_batched as jax_b
    from omnibias.struct.torch import soft_viterbi_batched as torch_b

    trellises = [random_chain(s) for s in range(4)]
    em = np.stack([t.emissions for t in trellises])
    trn = trellises[0].transitions
    st = trellises[0].start
    bt = torch_b(torch.tensor(em), torch.tensor(trn), 5.0, start=torch.tensor(st)).numpy()
    bj = np.asarray(jax_b(jnp.asarray(em), jnp.asarray(trn), 5.0, start=jnp.asarray(st)))
    assert np.max(np.abs(bt - bj)) < ATOL

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

Read the implementation

def soft_viterbi_batched(
    emissions: Array,
    transitions: Array,
    beta: float = 1.0,
    *,
    start: Array | None = None,
) -> Array:
    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 ``jax.vmap`` -- bit-identical to looping.
    """
    import jax

    t_dim = 0 if transitions.ndim == 3 else None
    if start is None:
        out: Array = jax.vmap(
            lambda e, tr: soft_viterbi(e, tr, beta), in_axes=(0, t_dim)
        )(emissions, transitions)
        return out
    s_dim = 0 if start.ndim == 2 else None
    out_s: Array = jax.vmap(
        lambda e, tr, st: soft_viterbi(e, tr, beta, start=st), in_axes=(0, t_dim, s_dim)
    )(emissions, transitions, start)
    return out_s

Related APIsomnibias.struct.jax.soft_viterbi↗

Other public import paths

• omnibias.struct.jax.soft_viterbi_batched
• omnibias.struct.jax.soft_dp.soft_viterbi_batched

FUNCTION

soft_viterbi_marginals

omnibias.struct.jax.soft_viterbi_marginalsSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

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[jax-viterbi]" -q

Regression example · structured attention parity and matches 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_structured_attention_parity_and_matches_marginals() -> None:
    rng = np.random.default_rng(5)
    scores = rng.standard_normal((5, 3))
    trans = rng.standard_normal((3, 3))
    beta = 1.5
    a_t = tattn(torch.tensor(scores), torch.tensor(trans), beta).numpy()
    a_j = np.asarray(jattn(jnp.asarray(scores), jnp.asarray(trans), beta))
    assert np.max(np.abs(a_t - a_j)) < 1e-9
    # It *is* the linear-chain marginal, by construction.
    marg = np.asarray(jsvm(jnp.asarray(scores), jnp.asarray(trans), beta))
    assert np.max(np.abs(a_j - marg)) < 1e-9
    _ = jsoftmax  # jax softmax import kept for symmetry with the torch oracle

python -m pytest packages/omnibias-struct/tests/test_plan_attn.py -k test_structured_attention_parity_and_matches_marginals -q

Read the implementation

def soft_viterbi_marginals(
    emissions: Array,
    transitions: Array,
    beta: float = 1.0,
    *,
    start: Array | None = None,
) -> Array:
    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 jnp.exp(beta * (alpha + bwd - value))

Related APIsomnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.soft_viterbi_marginals
• omnibias.struct.jax.soft_dp.soft_viterbi_marginals

FUNCTION

soft_viterbi_transition_marginals

omnibias.struct.jax.soft_viterbi_transition_marginalsSource ↗

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

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	ArrayRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · transition marginals equal grad jax

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_grad_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import soft_viterbi, soft_viterbi_transition_marginals

    trellis = random_chain(1)
    emissions = jnp.asarray(trellis.emissions)
    transitions = jnp.asarray(trellis.transitions)
    start = jnp.asarray(trellis.start)
    grad = jax.grad(lambda a: soft_viterbi(emissions, a, 6.0, start=start))(transitions)
    xi = soft_viterbi_transition_marginals(emissions, transitions, 6.0, start=start)
    assert np.allclose(np.asarray(xi).sum(axis=0), np.asarray(grad), atol=1e-9)

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

Regression example · transition marginals equal autograd 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_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 -k test_transition_marginals_equal_autograd_torch -q

Read the implementation

def soft_viterbi_transition_marginals(
    emissions: Array,
    transitions: Array,
    beta: float = 1.0,
    *,
    start: Array | None = None,
) -> Array:
    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[Array] = []
    for t in range(1, n_steps):
        term = alpha[t - 1][:, None] + transitions + (emissions[t] + bwd[t])[None, :]
        xis.append(jnp.exp(beta * (term - value)))
    return jnp.stack(xis, axis=0)  # (T - 1, S, S)

Related APIsomnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.soft_viterbi_transition_marginals
• omnibias.struct.jax.soft_dp.soft_viterbi_transition_marginals

← Package overviewRunnable learning examples →