API REFERENCE/omnibias-struct

omnibias.struct.jax.semiring

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.semiring

Module contract & mathematical conventions

Differentiable semiring / hypergraph DP driver (jax).

Bit-identical twin of omnibias.struct.torch.semiring (float64 -- enable jax_enable_x64). Runs the ~omnibias.struct._core.semiring.LogSemiring (lse_beta) reduction of a ~omnibias.struct._core.semiring.Hypergraph over a backend edge_weights vector (semiring_value) and reads off the closed-form edge marginals by an inside-outside sweep and the tower softmax (semiring_marginals); the marginal of edge e equals d value / d edge_weights[e]. Reproduces the hand-written soft-DP layers to < 1e-12 -- the additive-safety cross-check that the driver subsumes them without changing their numerics.

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

semiring_marginals

omnibias.struct.jax.semiring_marginalsSource ↗

semiring_marginals(graph: Hypergraph, edge_weights: Array, beta: float=1.0) -> Array

Closed-form edge marginals mu[e] = P_beta(edge e in the derivation) ((num_edges,)).

Inside-outside with the tower softmax; mu[e] = d semiring_value / d edge_weights[e] (the exact gradient), and the marginals of the edges into root sum to 1.

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	ArrayRequired	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 edge_weights.shape[0] != graph.num_edges, raises ValueError(f'edge_weights must have length num_edges={graph.num_edges}, got {edge_weights.shape[0]}').

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

Executed regression · sampler empirical marginals converge

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_sampler_empirical_marginals_converge(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_sampling(0, BETA, backend).ok

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

Regression example · driver 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_driver_torch_jax_parity(seed: int) -> None:
    dag = random_dag(seed)
    w = dag_weight_matrix(dag)
    g, idx = from_dag(dag)
    for beta in (1.0, 8.0):
        ew_t = _dag_edge_weights(g, idx, -torch.tensor(w), torch)
        ew_j = _dag_edge_weights(g, idx, -jnp.asarray(w), jnp)
        assert abs(float(st.semiring_value(g, ew_t, beta)) - float(sj.semiring_value(g, ew_j, beta))) < 1e-11
        mt = np.asarray(st.semiring_marginals(g, ew_t, beta).detach())
        mj = np.asarray(sj.semiring_marginals(g, ew_j, beta))
        assert float(np.max(np.abs(mt - mj))) < 1e-11

python -m pytest packages/omnibias-struct/tests/test_semiring.py -k test_driver_torch_jax_parity -q

Read the implementation

def semiring_marginals(graph: Hypergraph, edge_weights: Array, beta: float = 1.0) -> Array:
    r"""Closed-form edge marginals ``mu[e] = P_beta(edge e in the derivation)`` (``(num_edges,)``).

    Inside-outside with the tower softmax; ``mu[e] = d semiring_value / d edge_weights[e]``
    (the exact gradient), and the marginals of the edges into ``root`` sum to ``1``.
    """
    if edge_weights.shape[0] != graph.num_edges:
        raise ValueError(
            f"edge_weights must have length num_edges={graph.num_edges}, got {edge_weights.shape[0]}"
        )
    inside = _inside(graph, edge_weights, beta)
    value = inside[graph.root]
    neg = jnp.asarray(_NEG, dtype=edge_weights.dtype)
    outside: list[Array] = [neg] * graph.num_nodes
    outside[graph.root] = jnp.zeros((), dtype=edge_weights.dtype)
    as_tail: list[list[tuple[int, int]]] = [[] for _ in range(graph.num_nodes)]
    for ei, e in enumerate(graph.edges):
        for pos, t in enumerate(e.tails):
            as_tail[t].append((ei, pos))
    for v in range(graph.num_nodes - 1, -1, -1):
        if v == graph.root:
            continue
        rows: list[Array] = []
        for ei, pos in as_tail[v]:
            e = graph.edges[ei]
            term = edge_weights[ei] + outside[e.head]
            for j, t in enumerate(e.tails):
                if j != pos:
                    term = term + inside[t]
            rows.append(term)
        if rows:
            outside[v] = logsumexp_beta(jnp.stack(rows), beta, axis=-1)
    marg: list[Array] = []
    for ei, e in enumerate(graph.edges):
        term = edge_weights[ei] + outside[e.head]
        for t in e.tails:
            term = term + inside[t]
        marg.append(jnp.exp(beta * (term - value)))
    return jnp.stack(marg)

Related APIsomnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.semiring_marginals
• omnibias.struct.jax.semiring.semiring_marginals

FUNCTION

semiring_value

omnibias.struct.jax.semiring_valueSource ↗

semiring_value(graph: Hypergraph, edge_weights: Array, beta: float=1.0) -> Array

Soft (lse_beta) value of graph at its root for the edge_weights vector.

edge_weights is (graph.num_edges,); differentiable in it. Generalises omnibias.struct.jax.soft_shortest_path's topological reduce to arbitrary arity-1/2 hyperedges. -> max derivation score as beta -> inf.

Parameters

Name	Type & default	Contract
graphpositional-or-keyword	HypergraphRequired	See the operation contract and implementation below.
edge_weightspositional-or-keyword	ArrayRequired	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 edge_weights.shape[0] != graph.num_edges, raises ValueError(f'edge_weights must have length num_edges={graph.num_edges}, got {edge_weights.shape[0]}').

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

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 · driver reproduces dtw

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_driver_reproduces_dtw(seed: int) -> None:
    rng = np.random.default_rng(seed)
    cost = rng.standard_normal((4, 5)) ** 2
    g, specs = _dtw_graph(4, 5)
    for beta in (1.0, 8.0, 64.0):
        ew_t = torch.stack([-torch.tensor(cost)[i, j] for (i, j) in specs])
        assert abs(-float(st.semiring_value(g, ew_t, beta)) - float(st.soft_dtw(torch.tensor(cost), beta))) < TOL
        ew_j = jnp.stack([-jnp.asarray(cost)[i, j] for (i, j) in specs])
        assert abs(-float(sj.semiring_value(g, ew_j, beta)) - float(sj.soft_dtw(jnp.asarray(cost), beta))) < TOL

python -m pytest packages/omnibias-struct/tests/test_semiring.py -k test_driver_reproduces_dtw -q

Read the implementation

def semiring_value(graph: Hypergraph, edge_weights: Array, beta: float = 1.0) -> Array:
    r"""Soft (``lse_beta``) value of ``graph`` at its ``root`` for the ``edge_weights`` vector.

    ``edge_weights`` is ``(graph.num_edges,)``; differentiable in it. Generalises
    :func:`omnibias.struct.jax.soft_shortest_path`'s topological reduce to arbitrary
    arity-1/2 hyperedges. ``-> max derivation score`` as ``beta -> inf``.
    """
    if edge_weights.shape[0] != graph.num_edges:
        raise ValueError(
            f"edge_weights must have length num_edges={graph.num_edges}, got {edge_weights.shape[0]}"
        )
    value: Array = _inside(graph, edge_weights, beta)[graph.root]
    return value

Other public import paths

• omnibias.struct.jax.semiring_value
• omnibias.struct.jax.semiring.semiring_value

← Package overviewRunnable learning examples →