API REFERENCE/omnibias-struct

omnibias.struct.torch.parse

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.torch.parse

Module contract & mathematical conventions

Differentiable CKY inside / inside-outside parsing on the semiring driver (torch).

Bit-identical twin of omnibias.struct.jax.parse (float64). Lifts a ~omnibias.struct._core.parse.BinaryGrammar onto the shared hypergraph driver (omnibias.struct.torch.semiring_value / semiring_marginals) so parsing needs no bespoke recursion: soft_inside is the lse_beta inside partition over all parse trees (-> best-parse score as beta -> inf), and inside_outside reads off the closed-form **span / rule marginals** -- the exact gradient of soft_inside w.r.t. the emission and rule scores, which the tests pin equal to autograd.

The beta -> inf reduction is the relaxation axis; the delta -> 0 tower differentiates it exactly -- the same two axes as every other layer here.

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

inside_outside

omnibias.struct.torch.inside_outsideSource ↗

inside_outside(grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float=1.0) -> tuple[Tensor, Tensor]

Closed-form (emit_marginals, rule_marginals) -- the gradient of soft_inside.

emit_marginals (L, R) is the expected lexical usage (d soft_inside / d emit) and rule_marginals (num_rules,) the expected rule usage. Inside-outside via the tower softmax; equal to autograd of soft_inside.

Parameters

Name	Type & default	Contract
grammarpositional-or-keyword	BinaryGrammarRequired	See the operation contract and implementation below.
emitpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rulepositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor]

Executed regression · parse 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_parse_marginals_match_autodiff(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_parse_marginals(0, BETA, backend).ok

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

Regression example · inside outside equals autograd

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_inside_outside_equals_autograd(seed: int) -> None:
    emit, rule = _scores(seed, 4)
    et = torch.tensor(emit, requires_grad=True)
    rt = torch.tensor(rule, requires_grad=True)
    beta = 6.0
    value = st.soft_inside(GRAMMAR, et, rt, beta)
    em, rm = st.inside_outside(GRAMMAR, et, rt, beta)
    ge, gr = torch.autograd.grad(value, (et, rt))
    assert float((em.detach() - ge).abs().max()) < 1e-9
    assert float((rm.detach() - gr).abs().max()) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_parse.py -k test_inside_outside_equals_autograd -q

Read the implementation

def inside_outside(
    grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float = 1.0
) -> tuple[Tensor, Tensor]:
    r"""Closed-form ``(emit_marginals, rule_marginals)`` -- the gradient of :func:`soft_inside`.

    ``emit_marginals`` ``(L, R)`` is the expected lexical usage (``d soft_inside / d emit``)
    and ``rule_marginals`` ``(num_rules,)`` the expected rule usage. Inside-outside via the
    tower softmax; equal to ``autograd`` of :func:`soft_inside`.
    """
    spec = build_chart(grammar, int(emit.shape[0]))
    mu = semiring_marginals(spec.graph, _edge_weights(spec, emit, rule), beta)
    emit_marg = torch.zeros_like(emit)
    rule_marg = torch.zeros_like(rule)
    for e, s in enumerate(spec.edge_specs):
        if s[0] == "emit":
            emit_marg[int(s[1]), int(s[2])] = emit_marg[int(s[1]), int(s[2])] + mu[e]
        else:
            rule_marg[int(s[1])] = rule_marg[int(s[1])] + mu[e]
    return emit_marg, rule_marg

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

Other public import paths

• omnibias.struct.torch.inside_outside
• omnibias.struct.torch.parse.inside_outside

FUNCTION

soft_inside

omnibias.struct.torch.soft_insideSource ↗

soft_inside(grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float=1.0) -> Tensor

Soft inside partition beta^-1 log sum_trees exp(beta score) of a CKY chart.

emit is (L, R) lexical scores and rule is (num_rules,) binary-rule scores (both learnable). Differentiable in both; -> best-parse score as beta -> inf. The score of a tree is the sum of its lexical + rule scores.

Parameters

Name	Type & default	Contract
grammarpositional-or-keyword	BinaryGrammarRequired	See the operation contract and implementation below.
emitpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rulepositional-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 · parse 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_parse_marginals_match_autodiff(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_parse_marginals(0, BETA, backend).ok

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

Regression example · soft inside matches bruteforce

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_inside_matches_bruteforce(seed: int) -> None:
    emit, rule = _scores(seed, 4)
    et, rt = torch.tensor(emit), torch.tensor(rule)
    ej, rj = jnp.asarray(emit), jnp.asarray(rule)
    for beta in (1.0, 8.0, 64.0):
        ref = brute_force_cky(GRAMMAR, emit, rule, beta)
        assert abs(float(st.soft_inside(GRAMMAR, et, rt, beta)) - ref) < 1e-11
        assert abs(float(sj.soft_inside(GRAMMAR, ej, rj, beta)) - ref) < 1e-11

python -m pytest packages/omnibias-struct/tests/test_parse.py -k test_soft_inside_matches_bruteforce -q

Read the implementation

def soft_inside(grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float = 1.0) -> Tensor:
    r"""Soft inside partition ``beta^-1 log sum_trees exp(beta score)`` of a CKY chart.

    ``emit`` is ``(L, R)`` lexical scores and ``rule`` is ``(num_rules,)`` binary-rule
    scores (both learnable). Differentiable in both; ``-> best-parse score`` as
    ``beta -> inf``. The score of a tree is the sum of its lexical + rule scores.
    """
    spec = build_chart(grammar, int(emit.shape[0]))
    return semiring_value(spec.graph, _edge_weights(spec, emit, rule), beta)

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

Other public import paths

• omnibias.struct.torch.soft_inside
• omnibias.struct.torch.parse.soft_inside

FUNCTION

soft_inside_batched

omnibias.struct.torch.soft_inside_batchedSource ↗

soft_inside_batched(grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float=1.0) -> Tensor

Batched soft_inside -> (B,) for emit (B, L, R), rule (B, num_rules).

Shared grammar (so the chart hypergraph is identical across the batch); maps the inside partition over the leading batch axis with torch.func.vmap -- bit-identical to looping soft_inside.

Parameters

Name	Type & default	Contract
grammarpositional-or-keyword	BinaryGrammarRequired	See the operation contract and implementation below.
emitpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rulepositional-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 · soft inside and eisner 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_inside_and_eisner_batched_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import (
        soft_eisner,
        soft_eisner_batched,
        soft_inside,
        soft_inside_batched,
    )

    rng = np.random.default_rng(3)
    grammar = _grammar()
    emit = torch.tensor(rng.standard_normal((4, 3, 2)))
    rule = torch.tensor(rng.standard_normal((4, grammar.num_rules)))
    b_in = soft_inside_batched(grammar, emit, rule, 2.0)
    loop_in = torch.stack([soft_inside(grammar, emit[i], rule[i], 2.0) for i in range(4)])
    assert b_in.shape == (4,)
    assert torch.max(torch.abs(b_in - loop_in)).item() < ATOL

    arc = torch.tensor(rng.standard_normal((4, 4, 4)))
    b_ei = soft_eisner_batched(arc, 2.0)
    loop_ei = torch.stack([soft_eisner(arc[i], 2.0) for i in range(4)])
    assert torch.max(torch.abs(b_ei - loop_ei)).item() < ATOL

python -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_soft_inside_and_eisner_batched_torch" -q

Read the implementation

def soft_inside_batched(
    grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float = 1.0
) -> Tensor:
    r"""Batched :func:`soft_inside` -> ``(B,)`` for ``emit`` ``(B, L, R)``, ``rule`` ``(B, num_rules)``.

    Shared ``grammar`` (so the chart hypergraph is identical across the batch); maps the
    inside partition over the leading batch axis with ``torch.func.vmap`` -- bit-identical to
    looping :func:`soft_inside`.
    """
    from torch.func import vmap

    def fwd(e: Tensor, r: Tensor) -> Tensor:
        return soft_inside(grammar, e, r, beta)

    out: Tensor = vmap(fwd)(emit, rule)
    return out

Related APIsomnibias.struct.torch.soft_inside↗

Other public import paths

• omnibias.struct.torch.soft_inside_batched
• omnibias.struct.torch.parse.soft_inside_batched

FUNCTION

span_marginals

omnibias.struct.torch.span_marginalsSource ↗

span_marginals(grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float=1.0) -> Tensor

Span marginals P_beta(nonterminal A spans [i, j)) as an (R, L + 1, L + 1) tensor.

[A, i, j] is the probability that item (A, i, j) is used by a parse; the start symbol over the whole sentence has marginal 1. Summing the edge marginals into each chart node.

Parameters

Name	Type & default	Contract
grammarpositional-or-keyword	BinaryGrammarRequired	See the operation contract and implementation below.
emitpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rulepositional-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 · span marginals root is one

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_span_marginals_root_is_one() -> None:
    emit, rule = _scores(0, 5)
    sp = st.span_marginals(GRAMMAR, torch.tensor(emit), torch.tensor(rule), 4.0)
    assert abs(float(sp[GRAMMAR.start, 0, 5]) - 1.0) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_parse.py::test_span_marginals_root_is_one" -q

Read the implementation

def span_marginals(
    grammar: BinaryGrammar, emit: Tensor, rule: Tensor, beta: float = 1.0
) -> Tensor:
    r"""Span marginals ``P_beta(nonterminal A spans [i, j))`` as an ``(R, L + 1, L + 1)`` tensor.

    ``[A, i, j]`` is the probability that item ``(A, i, j)`` is used by a parse; the start
    symbol over the whole sentence has marginal ``1``. Summing the edge marginals into each
    chart node.
    """
    spec = build_chart(grammar, int(emit.shape[0]))
    mu = semiring_marginals(spec.graph, _edge_weights(spec, emit, rule), beta)
    length_l, r_nt = int(emit.shape[0]), grammar.num_nonterminals
    out = torch.zeros((r_nt, length_l + 1, length_l + 1), dtype=emit.dtype, device=emit.device)
    node_span = {node: span for span, node in spec.node_of.items()}
    for e in range(spec.graph.num_edges):
        a, i, j = node_span[spec.graph.edges[e].head]
        out[a, i, j] = out[a, i, j] + mu[e]
    return out

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

Other public import paths

• omnibias.struct.torch.span_marginals
• omnibias.struct.torch.parse.span_marginals

← Package overviewRunnable learning examples →