API REFERENCE/omnibias-struct

omnibias.struct.jax.parse

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.parse

Module contract & mathematical conventions

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

Bit-identical twin of omnibias.struct.torch.parse (float64 -- enable jax_enable_x64). Lifts a ~omnibias.struct._core.parse.BinaryGrammar onto the shared hypergraph driver: soft_inside is the lse_beta inside partition over all parse trees and inside_outside reads off the closed-form span / rule marginals (the exact gradient of soft_inside, pinned equal to autograd).

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.jax.inside_outsideSource ↗

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

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

Parameters

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

Returnstuple[Array, Array]

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

Regression example · 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_torch_jax_parity(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):
        assert abs(float(st.soft_inside(GRAMMAR, et, rt, beta)) - float(sj.soft_inside(GRAMMAR, ej, rj, beta))) < 1e-11
        emt, rmt = st.inside_outside(GRAMMAR, et, rt, beta)
        emj, rmj = sj.inside_outside(GRAMMAR, ej, rj, beta)
        assert float(np.max(np.abs(np.asarray(emt) - np.asarray(emj)))) < 1e-11
        assert float(np.max(np.abs(np.asarray(rmt) - np.asarray(rmj)))) < 1e-11

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

Read the implementation

def inside_outside(
    grammar: BinaryGrammar, emit: Array, rule: Array, beta: float = 1.0
) -> tuple[Array, Array]:
    r"""Closed-form ``(emit_marginals, rule_marginals)`` -- the gradient 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 = jnp.zeros_like(emit)
    rule_marg = jnp.zeros_like(rule)
    for e, s in enumerate(spec.edge_specs):
        if s[0] == "emit":
            emit_marg = emit_marg.at[int(s[1]), int(s[2])].add(mu[e])
        else:
            rule_marg = rule_marg.at[int(s[1])].add(mu[e])
    return emit_marg, rule_marg

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

Other public import paths

• omnibias.struct.jax.inside_outside
• omnibias.struct.jax.parse.inside_outside

FUNCTION

soft_inside

omnibias.struct.jax.soft_insideSource ↗

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

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

emit is (L, R) lexical scores and rule (num_rules,) binary-rule scores. Differentiable in both; -> best-parse score as beta -> inf.

Parameters

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

ReturnsArray

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[jax]" -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: Array, rule: Array, beta: float = 1.0) -> Array:
    r"""Soft inside partition ``beta^-1 log sum_trees exp(beta score)`` of a CKY chart.

    ``emit`` is ``(L, R)`` lexical scores and ``rule`` ``(num_rules,)`` binary-rule scores.
    Differentiable in both; ``-> best-parse score`` as ``beta -> inf``.
    """
    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.jax.semiring_value↗

Other public import paths

• omnibias.struct.jax.soft_inside
• omnibias.struct.jax.parse.soft_inside

FUNCTION

soft_inside_batched

omnibias.struct.jax.soft_inside_batchedSource ↗

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

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

Parameters

Name	Type & default	Contract
grammarpositional-or-keyword	BinaryGrammarRequired	See the operation contract and implementation below.
emitpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
rulepositional-or-keyword	ArrayRequired	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_inside_batched(
    grammar: BinaryGrammar, emit: Array, rule: Array, beta: float = 1.0
) -> Array:
    r"""Batched :func:`soft_inside` -> ``(B,)`` for ``emit`` ``(B, L, R)``, ``rule`` ``(B, num_rules)`` (``jax.vmap``)."""
    import jax

    out: Array = jax.vmap(lambda e, r: soft_inside(grammar, e, r, beta))(emit, rule)
    return out

Related APIsomnibias.struct.jax.soft_inside↗

Other public import paths

• omnibias.struct.jax.soft_inside_batched
• omnibias.struct.jax.parse.soft_inside_batched

FUNCTION

span_marginals

omnibias.struct.jax.span_marginalsSource ↗

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

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

Parameters

Name	Type & default	Contract
grammarpositional-or-keyword	BinaryGrammarRequired	See the operation contract and implementation below.
emitpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
rulepositional-or-keyword	ArrayRequired	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 span_marginals(grammar: BinaryGrammar, emit: Array, rule: Array, beta: float = 1.0) -> Array:
    r"""Span marginals ``P_beta(nonterminal A spans [i, j))`` as an ``(R, L + 1, L + 1)`` array."""
    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 = jnp.zeros((r_nt, length_l + 1, length_l + 1), dtype=emit.dtype)
    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 = out.at[a, i, j].add(mu[e])
    return out

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

Other public import paths

• omnibias.struct.jax.span_marginals
• omnibias.struct.jax.parse.span_marginals

← Package overviewRunnable learning examples →