API REFERENCE/omnibias-struct

omnibias.struct.torch.jets

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.torch.jets

Module contract & mathematical conventions

Exact higher-order jets of a soft-DP value through the unrolled recursion (torch).

This is the delta -> 0 founding bias collapse tower differentiating the beta -> inf relaxation to **all** orders (not just the first-order marginal). A univariate Taylor jet is propagated along a perturbation direction through every lse_beta reduction of the DP: each pairwise combine lse_beta(a, b) = a + beta^-1 softplus(beta (b - a)) has its whole jet from the beta-tempered softplus tower (softplus^(n) = sigma^(n-1) in omnibias.core) via omnibias.torch.jet.compose_jet, and multi-way reductions fold the exact associative pairwise jet. No autodiff and no finite differences -- the jet coefficients are closed form.

The returned jet is (order + 1,) Taylor coefficients c_k = f^(k)(0) / k! of t -> V(params + t * direction): c_0 is the value, c_1 is the directional gradient grad . direction, and 2 c_2 = direction^T H direction is the directional curvature. Bit-identical twin of omnibias.struct.jax.jets.

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

chain_lse_jet

omnibias.struct.torch.chain_lse_jetSource ↗

chain_lse_jet(emissions: Tensor, transitions: Tensor, direction: Tensor, beta: float=1.0, *, order: int=2, start: Tensor | None=None) -> Tensor

Exact jet of t -> soft_viterbi(emissions + t * direction, transitions, beta).

emissions and direction are (T, S); transitions is (S, S) (held constant); start is (S,) or None. Returns the (order + 1,) jet of the soft Viterbi value along direction in emission space. 2 * jet[2] is the directional curvature direction^T H direction of the value.

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.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderkeyword-only	intDefault: 2	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').

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

Executed regression · hvp matches closed form jet curvature

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_hvp_matches_closed_form_jet_curvature() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_directional_curvature, chain_lse_jet

    trellis = random_chain(0)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    d = torch.tensor(np.random.default_rng(1).standard_normal(e.shape))
    beta = 3.0
    curv_hvp = float(chain_directional_curvature(e, trn, d, beta, start=st))
    curv_jet = 2.0 * float(chain_lse_jet(e, trn, d, beta, order=2, start=st)[2])
    assert abs(curv_hvp - curv_jet) < ATOL

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_hvp_matches_closed_form_jet_curvature" -q

Regression example · chain jet matches finite differences order3 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_chain_jet_matches_finite_differences_order3_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi
    from omnibias.struct.torch.jets import chain_lse_jet

    trellis = random_chain(2)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    d = torch.tensor(np.random.default_rng(2).standard_normal(e.shape))
    beta = 2.0
    jet = chain_lse_jet(e, trn, d, beta, order=3, start=st)

    def value_at(t: float) -> float:
        return float(soft_viterbi(e + t * d, trn, beta, start=st))

    h = 1e-3
    fd1 = (value_at(h) - value_at(-h)) / (2 * h)
    fd2 = (value_at(h) - 2 * value_at(0.0) + value_at(-h)) / h**2
    assert abs(float(jet[1]) - fd1) < 1e-6
    assert abs(2.0 * float(jet[2]) - fd2) < 1e-4

python -m pytest packages/omnibias-struct/tests/test_jets.py -k test_chain_jet_matches_finite_differences_order3_torch -q

Read the implementation

def chain_lse_jet(
    emissions: Tensor,
    transitions: Tensor,
    direction: Tensor,
    beta: float = 1.0,
    *,
    order: int = 2,
    start: Tensor | None = None,
) -> Tensor:
    r"""Exact jet of ``t -> soft_viterbi(emissions + t * direction, transitions, beta)``.

    ``emissions`` and ``direction`` are ``(T, S)``; ``transitions`` is ``(S, S)`` (held
    constant); ``start`` is ``(S,)`` or ``None``. Returns the ``(order + 1,)`` jet of the
    soft Viterbi value along ``direction`` in emission space. ``2 * jet[2]`` is the
    directional curvature ``direction^T H direction`` of the value.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    n_steps, n_states = int(emissions.shape[0]), int(emissions.shape[1])
    start_t = torch.zeros(n_states, dtype=emissions.dtype) if start is None else start
    # alpha jet: (order + 1, S)
    alpha_cols = [
        _linear_jet(start_t[s] + emissions[0, s], direction[0, s], order) for s in range(n_states)
    ]
    alpha = torch.stack(alpha_cols, dim=1)  # (order + 1, S)
    for t in range(1, n_steps):
        new_cols = []
        for s in range(n_states):
            reduced = _lse_reduce_jet(
                [alpha[:, sp] + _const_jet(transitions[sp, s], order) for sp in range(n_states)],
                beta,
                order,
            )
            new_cols.append(reduced + _linear_jet(emissions[t, s], direction[t, s], order))
        alpha = torch.stack(new_cols, dim=1)
    return _lse_reduce_jet([alpha[:, s] for s in range(n_states)], beta, order)

Other public import paths

• omnibias.struct.torch.chain_lse_jet
• omnibias.struct.torch.jets.chain_lse_jet

FUNCTION

cky_lse_jet

omnibias.struct.torch.cky_lse_jetSource ↗

cky_lse_jet(grammar: BinaryGrammar, emit: Tensor, rule: Tensor, emit_dir: Tensor, rule_dir: Tensor, beta: float=1.0, *, order: int=2) -> Tensor

Exact jet of t -> soft_inside(emit + t emit_dir, rule + t rule_dir, beta).

emit / emit_dir are (L, R) and rule / rule_dir are (num_rules,); the returned (order + 1,) jet's c_1 is the directional gradient of the CKY inside partition and 2 c_2 its directional curvature -- pinned to backend autodiff of omnibias.struct.torch.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.
emit_dirpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rule_dirpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderkeyword-only	intDefault: 2	See the operation contract and implementation below.

ReturnsTensor

Executed regression · cky jet matches finite differences order3 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_cky_jet_matches_finite_differences_order3_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import soft_inside
    from omnibias.struct.torch.jets import cky_lse_jet

    grammar = _grammar()
    rng = np.random.default_rng(1)
    length = 3
    emit = torch.tensor(rng.standard_normal((length, grammar.num_nonterminals)))
    rule = torch.tensor(rng.standard_normal(grammar.num_rules))
    de = torch.tensor(rng.standard_normal(emit.shape))
    dr = torch.tensor(rng.standard_normal(rule.shape))
    beta = 2.0
    jet = cky_lse_jet(grammar, emit, rule, de, dr, beta, order=3)

    def value_at(t: float) -> float:
        return float(soft_inside(grammar, emit + t * de, rule + t * dr, beta))

    h = 1e-3
    fd1 = (value_at(h) - value_at(-h)) / (2 * h)
    fd2 = (value_at(h) - 2 * value_at(0.0) + value_at(-h)) / h**2
    assert abs(float(jet[1]) - fd1) < 1e-6
    assert abs(2.0 * float(jet[2]) - fd2) < 1e-4

python -m pytest "packages/omnibias-struct/tests/test_jets_families.py::test_cky_jet_matches_finite_differences_order3_torch" -q

Regression example · cky jet matches autodiff value grad hessian 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_cky_jet_matches_autodiff_value_grad_hessian_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import soft_inside
    from omnibias.struct.torch.jets import cky_lse_jet

    grammar = _grammar()
    rng = np.random.default_rng(0)
    length = 4
    emit = torch.tensor(rng.standard_normal((length, grammar.num_nonterminals)))
    rule = torch.tensor(rng.standard_normal(grammar.num_rules))
    de = torch.tensor(rng.standard_normal(emit.shape))
    dr = torch.tensor(rng.standard_normal(rule.shape))
    beta = 3.0
    jet = cky_lse_jet(grammar, emit, rule, de, dr, beta, order=2)

    e2 = emit.clone().requires_grad_(True)
    r2 = rule.clone().requires_grad_(True)
    value = soft_inside(grammar, e2, r2, beta)
    ge, gr = torch.autograd.grad(value, (e2, r2), create_graph=True)
    grad_d = (ge * de).sum() + (gr * dr).sum()
    hde, hdr = torch.autograd.grad(grad_d, (e2, r2), retain_graph=True)
    d_h_d = (hde * de).sum() + (hdr * dr).sum()

    assert abs(float(jet[0]) - float(value.detach())) < 1e-9
    assert abs(float(jet[1]) - float(grad_d.detach())) < 1e-9
    assert abs(2.0 * float(jet[2]) - float(d_h_d.detach())) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_jets_families.py -k test_cky_jet_matches_autodiff_value_grad_hessian_torch -q

Read the implementation

def cky_lse_jet(
    grammar: BinaryGrammar,
    emit: Tensor,
    rule: Tensor,
    emit_dir: Tensor,
    rule_dir: Tensor,
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Tensor:
    r"""Exact jet of ``t -> soft_inside(emit + t emit_dir, rule + t rule_dir, beta)``.

    ``emit`` / ``emit_dir`` are ``(L, R)`` and ``rule`` / ``rule_dir`` are ``(num_rules,)``;
    the returned ``(order + 1,)`` jet's ``c_1`` is the directional gradient of the CKY inside
    partition and ``2 c_2`` its directional curvature -- pinned to backend autodiff of
    :func:`omnibias.struct.torch.soft_inside`.
    """
    spec = build_chart(grammar, int(emit.shape[0]))
    weight_jets: list[Tensor] = []
    for s in spec.edge_specs:
        if s[0] == "emit":
            i, a = int(s[1]), int(s[2])
            weight_jets.append(_linear_jet(emit[i, a], emit_dir[i, a], order))
        else:
            r = int(s[1])
            weight_jets.append(_linear_jet(rule[r], rule_dir[r], order))
    return hypergraph_lse_jet(spec.graph, weight_jets, beta, order=order)

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

Other public import paths

• omnibias.struct.torch.cky_lse_jet
• omnibias.struct.torch.jets.cky_lse_jet

FUNCTION

dag_lse_jet

omnibias.struct.torch.dag_lse_jetSource ↗

dag_lse_jet(weights: Tensor, direction: Tensor, dag: DAG, beta: float=1.0, *, order: int=2) -> Tensor

Exact jet of t -> soft_shortest_path(weights + t * direction, dag, beta).

weights and direction are (n, n) edge tensors (only the dag edge entries matter). Returns the (order + 1,) jet of the softmin cost along direction. The score of edge (u, v) is -weight, so the value is -alpha[sink] in the max convention.

Parameters

Name	Type & default	Contract
weightspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
directionpositional-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.
orderkeyword-only	intDefault: 2	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').
• When sink 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 · dag jet matches autodiff 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_dag_jet_matches_autodiff_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_shortest_path
    from omnibias.struct.torch.jets import dag_lse_jet

    dag = random_dag(1)
    w = torch.tensor(dag_weight_matrix(dag))
    dw = torch.tensor(_dag_direction(dag, 1))
    beta = 3.0
    jet = dag_lse_jet(w, dw, dag, beta, order=2)

    w2 = w.clone().requires_grad_(True)
    cost = soft_shortest_path(w2, dag, beta)
    (grad,) = torch.autograd.grad(cost, w2, create_graph=True)
    grad_d = (grad * dw).sum()
    (hd,) = torch.autograd.grad(grad_d, w2, retain_graph=True)
    d_h_d = (hd * dw).sum()

    assert abs(float(jet[0]) - float(cost.detach())) < 1e-9
    assert abs(float(jet[1]) - float(grad_d.detach())) < 1e-9
    assert abs(2.0 * float(jet[2]) - float(d_h_d.detach())) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_jets.py::test_dag_jet_matches_autodiff_torch" -q

Regression example · chain jet matches finite differences order3 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_chain_jet_matches_finite_differences_order3_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import soft_viterbi
    from omnibias.struct.torch.jets import chain_lse_jet

    trellis = random_chain(2)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    d = torch.tensor(np.random.default_rng(2).standard_normal(e.shape))
    beta = 2.0
    jet = chain_lse_jet(e, trn, d, beta, order=3, start=st)

    def value_at(t: float) -> float:
        return float(soft_viterbi(e + t * d, trn, beta, start=st))

    h = 1e-3
    fd1 = (value_at(h) - value_at(-h)) / (2 * h)
    fd2 = (value_at(h) - 2 * value_at(0.0) + value_at(-h)) / h**2
    assert abs(float(jet[1]) - fd1) < 1e-6
    assert abs(2.0 * float(jet[2]) - fd2) < 1e-4

python -m pytest packages/omnibias-struct/tests/test_jets.py -k test_chain_jet_matches_finite_differences_order3_torch -q

Read the implementation

def dag_lse_jet(
    weights: Tensor,
    direction: Tensor,
    dag: DAG,
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Tensor:
    r"""Exact jet of ``t -> soft_shortest_path(weights + t * direction, dag, beta)``.

    ``weights`` and ``direction`` are ``(n, n)`` edge tensors (only the ``dag`` edge
    entries matter). Returns the ``(order + 1,)`` jet of the softmin cost along
    ``direction``. The score of edge ``(u, v)`` is ``-weight``, so the value is
    ``-alpha[sink]`` in the max convention.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    alpha: list[Tensor | None] = [None] * dag.num_nodes
    alpha[dag.source] = _const_jet(torch.zeros((), dtype=weights.dtype), order)
    for v in range(dag.num_nodes):
        if v == dag.source:
            continue
        preds = [u for u in dag.incoming(v) if alpha[u] is not None]
        if not preds:
            continue
        contribs = [
            alpha[u] - _linear_jet(weights[u, v], direction[u, v], order)  # type: ignore[operator]
            for u in preds
        ]
        alpha[v] = _lse_reduce_jet(contribs, beta, order)
    sink = alpha[dag.sink]
    if sink is None:
        raise ValueError("sink is not reachable from source in this DAG")
    return -sink

Other public import paths

• omnibias.struct.torch.dag_lse_jet
• omnibias.struct.torch.jets.dag_lse_jet

FUNCTION

dp_value_jet

omnibias.struct.torch.dp_value_jetSource ↗

dp_value_jet(problem: ChainTrellis | DAG, direction: Tensor, beta: float=1.0, *, order: int=2, weights: Tensor | None=None) -> Tensor

Dispatch the exact DP-value jet by problem type.

• ChainTrellis -- pulls emissions / transitions / start from the problem; direction perturbs the emissions (T, S).
• DAG -- requires the finite edge weights (n, n); direction perturbs them.

Parameters

Name	Type & default	Contract
problempositional-or-keyword	ChainTrellis | DAGRequired	See the operation contract and implementation below.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderkeyword-only	intDefault: 2	See the operation contract and implementation below.
weightskeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When weights is None, raises ValueError('dp_value_jet on a DAG requires the edge `weights` (n, n)').

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

Read the implementation

def dp_value_jet(
    problem: ChainTrellis | DAG,
    direction: Tensor,
    beta: float = 1.0,
    *,
    order: int = 2,
    weights: Tensor | None = None,
) -> Tensor:
    r"""Dispatch the exact DP-value jet by problem type.

    * :class:`ChainTrellis` -- pulls ``emissions`` / ``transitions`` / ``start`` from the
      problem; ``direction`` perturbs the emissions ``(T, S)``.
    * :class:`DAG` -- requires the finite edge ``weights`` ``(n, n)``; ``direction``
      perturbs them.
    """
    if isinstance(problem, ChainTrellis):
        return chain_lse_jet(
            torch.as_tensor(problem.emissions),
            torch.as_tensor(problem.transitions),
            torch.as_tensor(direction),
            beta,
            order=order,
            start=torch.as_tensor(problem.start),
        )
    if weights is None:
        raise ValueError("dp_value_jet on a DAG requires the edge `weights` (n, n)")
    return dag_lse_jet(torch.as_tensor(weights), torch.as_tensor(direction), problem, beta, order=order)

Related APIsomnibias.struct.torch.chain_lse_jet↗omnibias.struct.torch.dag_lse_jet↗

Other public import paths

• omnibias.struct.torch.dp_value_jet
• omnibias.struct.torch.jets.dp_value_jet

FUNCTION

eisner_lse_jet

omnibias.struct.torch.eisner_lse_jetSource ↗

eisner_lse_jet(arc: Tensor, direction: Tensor, beta: float=1.0, *, order: int=2) -> Tensor

Exact jet of t -> soft_eisner(arc + t * direction, beta).

arc / direction are (n + 1, n + 1) arc-score matrices; the returned (order + 1,) jet's c_1 is the directional gradient of the projective-parse partition and 2 c_2 its directional curvature -- pinned to backend autodiff of omnibias.struct.torch.soft_eisner.

Parameters

Name	Type & default	Contract
arcpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderkeyword-only	intDefault: 2	See the operation contract and implementation below.

ReturnsTensor

Executed regression · eisner jet matches autodiff 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_eisner_jet_matches_autodiff_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import soft_eisner
    from omnibias.struct.torch.jets import eisner_lse_jet

    rng = np.random.default_rng(2)
    n = 3
    arc = torch.tensor(rng.standard_normal((n + 1, n + 1)))
    d = torch.tensor(rng.standard_normal(arc.shape))
    beta = 3.0
    jet = eisner_lse_jet(arc, d, beta, order=2)

    a2 = arc.clone().requires_grad_(True)
    value = soft_eisner(a2, beta)
    (grad,) = torch.autograd.grad(value, a2, create_graph=True)
    grad_d = (grad * d).sum()
    (hd,) = torch.autograd.grad(grad_d, a2, retain_graph=True)
    d_h_d = (hd * d).sum()

    assert abs(float(jet[0]) - float(value.detach())) < 1e-9
    assert abs(float(jet[1]) - float(grad_d.detach())) < 1e-9
    assert abs(2.0 * float(jet[2]) - float(d_h_d.detach())) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_jets_families.py::test_eisner_jet_matches_autodiff_torch" -q

Regression example · cky and eisner jets 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_cky_and_eisner_jets_torch_jax_parity() -> None:
    pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    import torch
    from omnibias.struct.jax.jets import cky_lse_jet as jax_cky
    from omnibias.struct.jax.jets import eisner_lse_jet as jax_eis
    from omnibias.struct.torch.jets import cky_lse_jet as torch_cky
    from omnibias.struct.torch.jets import eisner_lse_jet as torch_eis

    torch.set_default_dtype(torch.float64)
    grammar = _grammar()
    rng = np.random.default_rng(4)
    length = 4
    emit = rng.standard_normal((length, grammar.num_nonterminals))
    rule = rng.standard_normal(grammar.num_rules)
    de = rng.standard_normal(emit.shape)
    dr = rng.standard_normal(rule.shape)
    ct = torch_cky(
        grammar, torch.tensor(emit), torch.tensor(rule), torch.tensor(de), torch.tensor(dr), 4.0, order=3
    ).numpy()
    cj = np.asarray(
        jax_cky(
            grammar, jnp.asarray(emit), jnp.asarray(rule), jnp.asarray(de), jnp.asarray(dr), 4.0, order=3
        )
    )
    assert np.max(np.abs(ct - cj)) < 1e-9

    n = 3
    arc = rng.standard_normal((n + 1, n + 1))
    d = rng.standard_normal((n + 1, n + 1))
    et = torch_eis(torch.tensor(arc), torch.tensor(d), 3.0, order=2).numpy()
    ej = np.asarray(jax_eis(jnp.asarray(arc), jnp.asarray(d), 3.0, order=2))
    assert np.max(np.abs(et - ej)) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_jets_families.py -k test_cky_and_eisner_jets_torch_jax_parity -q

Read the implementation

def eisner_lse_jet(
    arc: Tensor,
    direction: Tensor,
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Tensor:
    r"""Exact jet of ``t -> soft_eisner(arc + t * direction, beta)``.

    ``arc`` / ``direction`` are ``(n + 1, n + 1)`` arc-score matrices; the returned
    ``(order + 1,)`` jet's ``c_1`` is the directional gradient of the projective-parse
    partition and ``2 c_2`` its directional curvature -- pinned to backend autodiff of
    :func:`omnibias.struct.torch.soft_eisner`.
    """
    spec: EisnerSpec = eisner_hypergraph(int(arc.shape[0]) - 1)
    zero = torch.zeros((), dtype=arc.dtype, device=arc.device)
    weight_jets: list[Tensor] = []
    for s in spec.edge_specs:
        if s[0] == "arc":
            h, m = int(s[1]), int(s[2])
            weight_jets.append(_linear_jet(arc[h, m], direction[h, m], order))
        else:
            weight_jets.append(_const_jet(zero, order))
    return hypergraph_lse_jet(spec.graph, weight_jets, beta, order=order)

Related APIsomnibias.struct.eisner_hypergraph↗omnibias.struct.torch.hypergraph_lse_jet↗

Other public import paths

• omnibias.struct.torch.eisner_lse_jet
• omnibias.struct.torch.jets.eisner_lse_jet

FUNCTION

hypergraph_lse_jet

omnibias.struct.torch.hypergraph_lse_jetSource ↗

hypergraph_lse_jet(graph: Hypergraph, weight_jets: Sequence[Tensor], beta: float=1.0, *, order: int=2) -> Tensor

Exact jet of the lse_beta inside value of a Hypergraph at its root.

weight_jets is a length-graph.num_edges sequence of (order + 1,) scalar jets, one per hyperedge (typically a _linear_jet of a perturbed edge weight). The inside recursion inside[v] = lse_beta_e (weight[e] + sum_tails inside) is a +/lse_beta fold, so every step is either a jet sum (_sum_jets) or the exact associative pairwise lse2_jet fold (_lse_reduce_jet) -- the whole jet is closed form, no autodiff. This is the driver the CKY / Eisner jets lift to (mirroring omnibias.struct.torch.semiring_value).

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
weight_jetspositional-or-keyword	Sequence[Tensor]Required	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderkeyword-only	intDefault: 2	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').
• When len(weight_jets) != graph.num_edges, raises ValueError(f'weight_jets must have length num_edges={graph.num_edges}, got {len(weight_jets)}').
• When root is None, raises ValueError('root is unreachable from the axioms').

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

Executed regression · eisner jet matches autodiff 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_eisner_jet_matches_autodiff_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import soft_eisner
    from omnibias.struct.torch.jets import eisner_lse_jet

    rng = np.random.default_rng(2)
    n = 3
    arc = torch.tensor(rng.standard_normal((n + 1, n + 1)))
    d = torch.tensor(rng.standard_normal(arc.shape))
    beta = 3.0
    jet = eisner_lse_jet(arc, d, beta, order=2)

    a2 = arc.clone().requires_grad_(True)
    value = soft_eisner(a2, beta)
    (grad,) = torch.autograd.grad(value, a2, create_graph=True)
    grad_d = (grad * d).sum()
    (hd,) = torch.autograd.grad(grad_d, a2, retain_graph=True)
    d_h_d = (hd * d).sum()

    assert abs(float(jet[0]) - float(value.detach())) < 1e-9
    assert abs(float(jet[1]) - float(grad_d.detach())) < 1e-9
    assert abs(2.0 * float(jet[2]) - float(d_h_d.detach())) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_jets_families.py::test_eisner_jet_matches_autodiff_torch" -q

Read the implementation

def hypergraph_lse_jet(
    graph: Hypergraph,
    weight_jets: Sequence[Tensor],
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Tensor:
    r"""Exact jet of the ``lse_beta`` inside value of a :class:`Hypergraph` at its ``root``.

    ``weight_jets`` is a length-``graph.num_edges`` sequence of ``(order + 1,)`` scalar jets,
    one per hyperedge (typically a :func:`_linear_jet` of a perturbed edge weight). The inside
    recursion ``inside[v] = lse_beta_e (weight[e] + sum_tails inside)`` is a
    ``+``/``lse_beta`` fold, so every step is either a jet sum (:func:`_sum_jets`) or the exact
    associative pairwise :func:`lse2_jet` fold (:func:`_lse_reduce_jet`) -- the whole jet is
    closed form, no autodiff. This is the driver the CKY / Eisner jets lift to (mirroring
    :func:`omnibias.struct.torch.semiring_value`).
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    if len(weight_jets) != graph.num_edges:
        raise ValueError(
            f"weight_jets must have length num_edges={graph.num_edges}, got {len(weight_jets)}"
        )
    node_jets: list[Tensor | None] = [None] * graph.num_nodes
    for v in range(graph.num_nodes):
        edge_jets: list[Tensor] = []
        for ei in graph.incoming(v):
            tail_jets = [node_jets[t] for t in graph.edges[ei].tails]
            if any(tj is None for tj in tail_jets):
                continue  # a derivation through an unreachable tail is the -inf semiring zero
            edge_jets.append(_sum_jets([weight_jets[ei], *tail_jets]))  # type: ignore[list-item]
        if edge_jets:
            node_jets[v] = _lse_reduce_jet(edge_jets, beta, order)
    root = node_jets[graph.root]
    if root is None:
        raise ValueError("root is unreachable from the axioms")
    return root

Other public import paths

• omnibias.struct.torch.hypergraph_lse_jet
• omnibias.struct.torch.jets.hypergraph_lse_jet

FUNCTION

lse2_jet

omnibias.struct.torch.lse2_jetSource ↗

lse2_jet(a: Tensor, b: Tensor, beta: float=1.0, order: int=2) -> Tensor

Jet of t -> lse_beta(a(t), b(t)) from two input jets a, b.

a and b are (order + 1, ...) Taylor-coefficient jets sharing a trailing shape. Uses lse_beta(a, b) = a + beta^-1 softplus(beta (b - a)) and composes the softplus tower through omnibias.torch.jet.compose_jet. Generalises omnibias.struct.torch.pairwise_lse_jet (which fixes a constant).

Parameters

Name	Type & default	Contract
apositional-or-keyword	TensorRequired	See the operation contract and implementation below.
bpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderpositional-or-keyword	intDefault: 2	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').

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

Executed regression · hvp matches closed form jet curvature

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_hvp_matches_closed_form_jet_curvature() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_directional_curvature, chain_lse_jet

    trellis = random_chain(0)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    d = torch.tensor(np.random.default_rng(1).standard_normal(e.shape))
    beta = 3.0
    curv_hvp = float(chain_directional_curvature(e, trn, d, beta, start=st))
    curv_jet = 2.0 * float(chain_lse_jet(e, trn, d, beta, order=2, start=st)[2])
    assert abs(curv_hvp - curv_jet) < ATOL

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_hvp_matches_closed_form_jet_curvature" -q

Regression example · lse2 jet generalizes pairwise lse jet 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_lse2_jet_generalizes_pairwise_lse_jet_torch() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import pairwise_lse_jet
    from omnibias.struct.torch.jets import lse2_jet

    a0 = torch.tensor(0.3)
    b0 = torch.tensor(-0.4)
    db = torch.tensor(1.1)
    beta = 2.0
    order = 3
    # pairwise_lse_jet fixes a constant and moves b linearly: build the same two jets.
    a_jet = torch.stack([a0] + [torch.zeros(()) for _ in range(order)])
    b_jet = torch.stack([b0, db] + [torch.zeros(()) for _ in range(order - 1)])
    got = lse2_jet(a_jet, b_jet, beta, order)
    ref = pairwise_lse_jet(a0, b0, db, beta, order)
    assert np.max(np.abs(got.numpy() - ref.numpy())) < 1e-12

python -m pytest packages/omnibias-struct/tests/test_jets.py -k test_lse2_jet_generalizes_pairwise_lse_jet_torch -q

Read the implementation

def lse2_jet(a: Tensor, b: Tensor, beta: float = 1.0, order: int = 2) -> Tensor:
    r"""Jet of ``t -> lse_beta(a(t), b(t))`` from two input jets ``a``, ``b``.

    ``a`` and ``b`` are ``(order + 1, ...)`` Taylor-coefficient jets sharing a trailing
    shape. Uses ``lse_beta(a, b) = a + beta^-1 softplus(beta (b - a))`` and composes the
    ``softplus`` tower through :func:`omnibias.torch.jet.compose_jet`. Generalises
    :func:`omnibias.struct.torch.pairwise_lse_jet` (which fixes ``a`` constant).
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    u = beta * (b - a)
    tower = _softplus_tower(u[0], order)
    sp = compose_jet(u, tower) / beta
    result: Tensor = a + sp
    return result

Related APIsomnibias.torch.compose_jet↗

Other public import paths

• omnibias.struct.torch.lse2_jet
• omnibias.struct.torch.jets.lse2_jet

← Package overviewRunnable learning examples →