API REFERENCE/omnibias-struct

omnibias.struct.jax.jets

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.jets

Module contract & mathematical conventions

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

Bit-identical twin of omnibias.struct.torch.jets (float64 -- enable jax_enable_x64). This is the delta -> 0 founding bias collapse tower differentiating the beta -> inf relaxation to **all** orders: 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)) carrying its whole jet from the beta-tempered softplus tower (softplus^(n) = sigma^(n-1)) via omnibias.jax.jet.compose_jet, and multi-way reductions folding the exact associative pairwise jet.

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

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

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

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.

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.
directionpositional-or-keyword	ArrayRequired	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	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

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 · chain and dag jets 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_chain_and_dag_jets_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.jets import chain_lse_jet as jax_chain
    from omnibias.struct.jax.jets import dag_lse_jet as jax_dag
    from omnibias.struct.torch.jets import chain_lse_jet as torch_chain
    from omnibias.struct.torch.jets import dag_lse_jet as torch_dag

    trellis = random_chain(3)
    d = np.random.default_rng(3).standard_normal(trellis.emissions.shape)
    jt = torch_chain(
        torch.tensor(trellis.emissions), torch.tensor(trellis.transitions), torch.tensor(d),
        4.0, order=3, start=torch.tensor(trellis.start),
    ).numpy()
    jj = np.asarray(
        jax_chain(
            jnp.asarray(trellis.emissions), jnp.asarray(trellis.transitions), jnp.asarray(d),
            4.0, order=3, start=jnp.asarray(trellis.start),
        )
    )
    assert np.max(np.abs(jt - jj)) < 1e-9

    dag = random_dag(2)
    w = dag_weight_matrix(dag)
    dw = _dag_direction(dag, 2)
    dt = torch_dag(torch.tensor(w), torch.tensor(dw), dag, 3.0, order=2).numpy()
    dj = np.asarray(jax_dag(jnp.asarray(w), jnp.asarray(dw), dag, 3.0, order=2))
    assert np.max(np.abs(dt - dj)) < 1e-9

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

Read the implementation

def chain_lse_jet(
    emissions: Array,
    transitions: Array,
    direction: Array,
    beta: float = 1.0,
    *,
    order: int = 2,
    start: Array | None = None,
) -> Array:
    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.
    """
    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 = jnp.zeros(n_states, dtype=emissions.dtype) if start is None else start
    alpha_cols = [
        _linear_jet(start_t[s] + emissions[0, s], direction[0, s], order) for s in range(n_states)
    ]
    alpha = jnp.stack(alpha_cols, axis=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 = jnp.stack(new_cols, axis=1)
    return _lse_reduce_jet([alpha[:, s] for s in range(n_states)], beta, order)

Other public import paths

• omnibias.struct.jax.chain_lse_jet
• omnibias.struct.jax.jets.chain_lse_jet

FUNCTION

cky_lse_jet

omnibias.struct.jax.cky_lse_jetSource ↗

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

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.jax.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.
emit_dirpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
rule_dirpositional-or-keyword	ArrayRequired	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.

ReturnsArray

Executed regression · cky and eisner jets 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_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::test_cky_and_eisner_jets_torch_jax_parity" -q

Read the implementation

def cky_lse_jet(
    grammar: BinaryGrammar,
    emit: Array,
    rule: Array,
    emit_dir: Array,
    rule_dir: Array,
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Array:
    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.jax.soft_inside`.
    """
    spec = build_chart(grammar, int(emit.shape[0]))
    weight_jets: list[Array] = []
    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.jax.hypergraph_lse_jet↗

Other public import paths

• omnibias.struct.jax.cky_lse_jet
• omnibias.struct.jax.jets.cky_lse_jet

FUNCTION

dag_lse_jet

omnibias.struct.jax.dag_lse_jetSource ↗

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

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 (value = -alpha[sink] in the max convention).

Parameters

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

ReturnsArray

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 · chain and dag jets 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_chain_and_dag_jets_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.jets import chain_lse_jet as jax_chain
    from omnibias.struct.jax.jets import dag_lse_jet as jax_dag
    from omnibias.struct.torch.jets import chain_lse_jet as torch_chain
    from omnibias.struct.torch.jets import dag_lse_jet as torch_dag

    trellis = random_chain(3)
    d = np.random.default_rng(3).standard_normal(trellis.emissions.shape)
    jt = torch_chain(
        torch.tensor(trellis.emissions), torch.tensor(trellis.transitions), torch.tensor(d),
        4.0, order=3, start=torch.tensor(trellis.start),
    ).numpy()
    jj = np.asarray(
        jax_chain(
            jnp.asarray(trellis.emissions), jnp.asarray(trellis.transitions), jnp.asarray(d),
            4.0, order=3, start=jnp.asarray(trellis.start),
        )
    )
    assert np.max(np.abs(jt - jj)) < 1e-9

    dag = random_dag(2)
    w = dag_weight_matrix(dag)
    dw = _dag_direction(dag, 2)
    dt = torch_dag(torch.tensor(w), torch.tensor(dw), dag, 3.0, order=2).numpy()
    dj = np.asarray(jax_dag(jnp.asarray(w), jnp.asarray(dw), dag, 3.0, order=2))
    assert np.max(np.abs(dt - dj)) < 1e-9

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

Read the implementation

def dag_lse_jet(
    weights: Array,
    direction: Array,
    dag: DAG,
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Array:
    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`` (``value = -alpha[sink]`` in the max convention).
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    alpha: list[Array | None] = [None] * dag.num_nodes
    alpha[dag.source] = _const_jet(jnp.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.jax.dag_lse_jet
• omnibias.struct.jax.jets.dag_lse_jet

FUNCTION

dp_value_jet

omnibias.struct.jax.dp_value_jetSource ↗

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

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	ArrayRequired	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	Array | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

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: Array,
    beta: float = 1.0,
    *,
    order: int = 2,
    weights: Array | None = None,
) -> Array:
    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(
            jnp.asarray(problem.emissions),
            jnp.asarray(problem.transitions),
            jnp.asarray(direction),
            beta,
            order=order,
            start=jnp.asarray(problem.start),
        )
    if weights is None:
        raise ValueError("dp_value_jet on a DAG requires the edge `weights` (n, n)")
    return dag_lse_jet(jnp.asarray(weights), jnp.asarray(direction), problem, beta, order=order)

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

Other public import paths

• omnibias.struct.jax.dp_value_jet
• omnibias.struct.jax.jets.dp_value_jet

FUNCTION

eisner_lse_jet

omnibias.struct.jax.eisner_lse_jetSource ↗

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

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.jax.soft_eisner.

Parameters

Name	Type & default	Contract
arcpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
directionpositional-or-keyword	ArrayRequired	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.

ReturnsArray

Executed regression · eisner jet value matches soft eisner 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_eisner_jet_value_matches_soft_eisner_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import soft_eisner
    from omnibias.struct.jax.jets import eisner_lse_jet

    rng = np.random.default_rng(3)
    n = 3
    arc = jnp.asarray(rng.standard_normal((n + 1, n + 1)))
    d = jnp.asarray(rng.standard_normal((n + 1, n + 1)))
    beta = 2.5
    jet = eisner_lse_jet(arc, d, beta, order=2)
    value = soft_eisner(arc, beta)
    grad = jax.grad(lambda a: soft_eisner(a, beta))(arc)
    grad_d = float(jnp.sum(grad * d))
    assert abs(float(jet[0]) - float(value)) < 1e-9
    assert abs(float(jet[1]) - grad_d) < 1e-9

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

Regression example · eisner jet matches autodiff 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_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 -k test_eisner_jet_matches_autodiff_torch -q

Read the implementation

def eisner_lse_jet(
    arc: Array,
    direction: Array,
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Array:
    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.jax.soft_eisner`.
    """
    spec: EisnerSpec = eisner_hypergraph(int(arc.shape[0]) - 1)
    zero = jnp.zeros((), dtype=arc.dtype)
    weight_jets: list[Array] = []
    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.jax.hypergraph_lse_jet↗

Other public import paths

• omnibias.struct.jax.eisner_lse_jet
• omnibias.struct.jax.jets.eisner_lse_jet

FUNCTION

hypergraph_lse_jet

omnibias.struct.jax.hypergraph_lse_jetSource ↗

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

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.jax.semiring_value).

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
weight_jetspositional-or-keyword	Sequence[Array]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.

ReturnsArray

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 value matches soft eisner 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_eisner_jet_value_matches_soft_eisner_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import soft_eisner
    from omnibias.struct.jax.jets import eisner_lse_jet

    rng = np.random.default_rng(3)
    n = 3
    arc = jnp.asarray(rng.standard_normal((n + 1, n + 1)))
    d = jnp.asarray(rng.standard_normal((n + 1, n + 1)))
    beta = 2.5
    jet = eisner_lse_jet(arc, d, beta, order=2)
    value = soft_eisner(arc, beta)
    grad = jax.grad(lambda a: soft_eisner(a, beta))(arc)
    grad_d = float(jnp.sum(grad * d))
    assert abs(float(jet[0]) - float(value)) < 1e-9
    assert abs(float(jet[1]) - grad_d) < 1e-9

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

Read the implementation

def hypergraph_lse_jet(
    graph: Hypergraph,
    weight_jets: Sequence[Array],
    beta: float = 1.0,
    *,
    order: int = 2,
) -> Array:
    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.jax.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[Array | None] = [None] * graph.num_nodes
    for v in range(graph.num_nodes):
        edge_jets: list[Array] = []
        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.jax.hypergraph_lse_jet
• omnibias.struct.jax.jets.hypergraph_lse_jet

FUNCTION

lse2_jet

omnibias.struct.jax.lse2_jetSource ↗

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

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.jax.jet.compose_jet. Generalises omnibias.struct.jax.pairwise_lse_jet (which fixes a constant).

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-or-keyword	ArrayRequired	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.

ReturnsArray

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 · eisner jet value matches soft eisner 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_eisner_jet_value_matches_soft_eisner_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import soft_eisner
    from omnibias.struct.jax.jets import eisner_lse_jet

    rng = np.random.default_rng(3)
    n = 3
    arc = jnp.asarray(rng.standard_normal((n + 1, n + 1)))
    d = jnp.asarray(rng.standard_normal((n + 1, n + 1)))
    beta = 2.5
    jet = eisner_lse_jet(arc, d, beta, order=2)
    value = soft_eisner(arc, beta)
    grad = jax.grad(lambda a: soft_eisner(a, beta))(arc)
    grad_d = float(jnp.sum(grad * d))
    assert abs(float(jet[0]) - float(value)) < 1e-9
    assert abs(float(jet[1]) - grad_d) < 1e-9

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

Read the implementation

def lse2_jet(a: Array, b: Array, beta: float = 1.0, order: int = 2) -> Array:
    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.jax.jet.compose_jet`. Generalises
    :func:`omnibias.struct.jax.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: Array = a + sp
    return result

Related APIsomnibias.jax.compose_jet↗

Other public import paths

• omnibias.struct.jax.lse2_jet
• omnibias.struct.jax.jets.lse2_jet

← Package overviewRunnable learning examples →