API REFERENCE/omnibias-struct

omnibias.struct.jax.distributions

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.distributions

Module contract & mathematical conventions

Distribution operators over the Gibbs path distribution p_beta (JAX).

Bit-identical twin of omnibias.struct.torch.distributions for the differentiable ops. For the Gibbs distribution p_beta(D) proportional to exp(beta * score(D)) over the derivations of a ~omnibias.struct._core.semiring.Hypergraph:

• path_entropy -- the exact Shannon entropy H(p_beta) = beta (V_beta - E_p[score]) in closed form from the inside value and the edge marginals (the first-order expectation-semiring quantity); differentiable, >= 0, -> log(#derivations) as beta -> 0 and -> log(#argmax) as beta -> inf.
• sample_paths -- **exact** forward-filtering backward-sampling from p_beta (not a relaxation); its empirical edge frequencies converge to semiring_marginals.
• gumbel_relaxed_sample -- a Gumbel-perturb-and-MAP **approximate**, differentiable relaxed sample (optionally straight-through); honestly labelled, *not* an exact draw.
• topk_paths -- **exact** k-best derivations (a hard decode) and topk_free_energy -- the differentiable lse_beta restricted to those top-k scores (-> best score as beta -> inf, -> the full soft value as k -> #derivations).

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

gumbel_relaxed_sample

omnibias.struct.jax.gumbel_relaxed_sampleSource ↗

gumbel_relaxed_sample(graph: Hypergraph, edge_weights: Array, beta: float=1.0, *, seed: int | None=None, hard: bool=False) -> Array

Gumbel perturb-and-MAP **approximate** relaxed sample (differentiable), as edge scores.

Adds i.i.d. Gumbel noise to the edge weights and returns the tower-softmax edge selection semiring_marginals of the perturbed weights; with hard=True a straight-through estimator returns the argmax-derivation indicator in the forward pass and the relaxed gradient in the backward. This is a per-edge perturbation, **not** an exact draw from p_beta (use sample_paths for that) -- honestly labelled as a relaxation.

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.
seedkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.
hardkeyword-only	boolDefault: False	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def gumbel_relaxed_sample(
    graph: Hypergraph,
    edge_weights: Array,
    beta: float = 1.0,
    *,
    seed: int | None = None,
    hard: bool = False,
) -> Array:
    r"""Gumbel perturb-and-MAP **approximate** relaxed sample (differentiable), as edge scores.

    Adds i.i.d. Gumbel noise to the edge weights and returns the tower-softmax edge selection
    :func:`semiring_marginals` of the perturbed weights; with ``hard=True`` a straight-through
    estimator returns the argmax-derivation indicator in the forward pass and the relaxed
    gradient in the backward. This is a per-edge perturbation, **not** an exact draw from
    ``p_beta`` (use :func:`sample_paths` for that) -- honestly labelled as a relaxation.
    """
    rng = np.random.default_rng(seed)
    u = rng.random(size=int(edge_weights.shape[0]))
    gumbel = jnp.asarray(-np.log(-np.log(u)), dtype=edge_weights.dtype)
    perturbed = edge_weights + gumbel
    relaxed: Array = semiring_marginals(graph, perturbed, beta)
    if not hard:
        return relaxed
    _score, deriv = best_derivation(graph, [float(x) for x in np.asarray(perturbed).reshape(-1)])
    hard_ind = jnp.zeros_like(edge_weights)
    for e in deriv:
        hard_ind = hard_ind.at[e].add(1.0)
    straight_through: Array = hard_ind + (relaxed - lax.stop_gradient(relaxed))
    return straight_through

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

Other public import paths

• omnibias.struct.jax.gumbel_relaxed_sample
• omnibias.struct.jax.distributions.gumbel_relaxed_sample

FUNCTION

path_entropy

omnibias.struct.jax.path_entropySource ↗

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

Exact Shannon entropy (nats) of p_beta over the derivations of graph.

Closed form H = beta (V_beta - E_p[score]) = beta (value - <marginals, weights>) -- the first-order expectation-semiring entropy. Differentiable in edge_weights; equals omnibias.struct._core.operators.brute_force_entropy on tiny graphs.

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

Executed regression · entropy identity and mc

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_entropy_identity_and_mc(backend: str) -> None:
    pytest.importorskip(backend)
    rec = probe_entropy(0, BETA, backend)
    assert rec.ok, rec
    assert rec.metrics["identity_err"] < 1e-8

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

Regression example · path entropy 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_path_entropy_matches_bruteforce(seed: int) -> None:
    w = _weights(seed)
    wt, wj = torch.tensor(w), jnp.asarray(w)
    for beta in (0.5, 2.0, 8.0):
        ref = brute_force_entropy(GRAPH, w, beta)
        assert abs(float(st.path_entropy(GRAPH, wt, beta)) - ref) < 1e-9
        assert abs(float(sj.path_entropy(GRAPH, wj, beta)) - ref) < 1e-9

python -m pytest packages/omnibias-struct/tests/test_operators.py -k test_path_entropy_matches_bruteforce -q

Read the implementation

def path_entropy(graph: Hypergraph, edge_weights: Array, beta: float = 1.0) -> Array:
    r"""Exact Shannon entropy (nats) of ``p_beta`` over the derivations of ``graph``.

    Closed form ``H = beta (V_beta - E_p[score]) = beta (value - <marginals, weights>)`` --
    the first-order expectation-semiring entropy. Differentiable in ``edge_weights``; equals
    :func:`omnibias.struct._core.operators.brute_force_entropy` on tiny graphs.
    """
    value = semiring_value(graph, edge_weights, beta)
    marginals = semiring_marginals(graph, edge_weights, beta)
    expected_score = (marginals * edge_weights).sum()
    entropy: Array = beta * (value - expected_score)
    return entropy

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

Other public import paths

• omnibias.struct.jax.path_entropy
• omnibias.struct.jax.distributions.path_entropy

FUNCTION

sample_paths

omnibias.struct.jax.sample_pathsSource ↗

sample_paths(graph: Hypergraph, edge_weights: Array, beta: float=1.0, num_samples: int=1, *, seed: int | None=None) -> tuple[Array, list[tuple[int, ...]]]

Exact forward-filtering backward-sampling of num_samples derivations from p_beta.

Returns (counts, samples): counts is a (num_samples, num_edges) array of per-edge usage counts and samples the chosen edge-index tuples. Exact (not a relaxation); counts.mean(0) -> semiring_marginals as num_samples -> inf. Sampling is not differentiable (use gumbel_relaxed_sample for a relaxed draw).

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.
num_samplespositional-or-keyword	intDefault: 1	See the operation contract and implementation below.
seedkeyword-only	int | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[Array, list[tuple[int, ...]]]

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

Read the implementation

def sample_paths(
    graph: Hypergraph,
    edge_weights: Array,
    beta: float = 1.0,
    num_samples: int = 1,
    *,
    seed: int | None = None,
) -> tuple[Array, list[tuple[int, ...]]]:
    r"""Exact forward-filtering backward-sampling of ``num_samples`` derivations from ``p_beta``.

    Returns ``(counts, samples)``: ``counts`` is a ``(num_samples, num_edges)`` array of
    per-edge usage counts and ``samples`` the chosen edge-index tuples. Exact (not a
    relaxation); ``counts.mean(0) -> `` :func:`semiring_marginals` as ``num_samples -> inf``.
    Sampling is not differentiable (use :func:`gumbel_relaxed_sample` for a relaxed draw).
    """
    weights = [float(x) for x in np.asarray(edge_weights).reshape(-1)]
    counts, samples = sample_derivations(graph, weights, beta, num_samples, seed=seed)
    return jnp.asarray(counts, dtype=edge_weights.dtype), samples

Related APIsomnibias.struct.sample_derivations↗

Other public import paths

• omnibias.struct.jax.sample_paths
• omnibias.struct.jax.distributions.sample_paths

FUNCTION

topk_free_energy

omnibias.struct.jax.topk_free_energySource ↗

topk_free_energy(graph: Hypergraph, edge_weights: Array, k: int, beta: float=1.0) -> Array

Differentiable lse_beta restricted to the top-k derivation scores.

The top-k *set* is chosen by exact hard k-best on the detached weights (piecewise constant in the weights); the returned free energy is differentiable through the k selected derivations' scores. -> the best score as beta -> inf and -> the full semiring_value as k -> #derivations (monotone non-decreasing in k).

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.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.

ReturnsArray

Executed regression · topk free energy 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_topk_free_energy_parity(seed: int) -> None:
    w = _weights(seed)
    wt, wj = torch.tensor(w), jnp.asarray(w)
    for k in (1, 3, 6):
        t = float(st.topk_free_energy(GRAPH, wt, k, 2.0))
        j = float(sj.topk_free_energy(GRAPH, wj, k, 2.0))
        assert abs(t - j) < 1e-10

python -m pytest "packages/omnibias-struct/tests/test_operators.py::test_topk_free_energy_parity[0]" -q

Read the implementation

def topk_free_energy(graph: Hypergraph, edge_weights: Array, k: int, beta: float = 1.0) -> Array:
    r"""Differentiable ``lse_beta`` restricted to the top-``k`` derivation scores.

    The top-``k`` *set* is chosen by exact hard k-best on the detached weights (piecewise
    constant in the weights); the returned free energy is differentiable through the ``k``
    selected derivations' scores. ``-> `` the best score as ``beta -> inf`` and ``-> `` the
    full :func:`semiring_value` as ``k -> #derivations`` (monotone non-decreasing in ``k``).
    """
    weights = [float(x) for x in np.asarray(edge_weights).reshape(-1)]
    kbest = kbest_derivations(graph, weights, k)
    scores = jnp.stack(
        [
            jnp.stack([edge_weights[e] for e in deriv]).sum() if deriv else jnp.zeros((), dtype=edge_weights.dtype)
            for _score, deriv in kbest
        ]
    )
    free_energy: Array = logsumexp_beta(scores, beta, axis=-1)
    return free_energy

Related APIsomnibias.struct.kbest_derivations↗omnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.topk_free_energy
• omnibias.struct.jax.distributions.topk_free_energy

FUNCTION

topk_paths

omnibias.struct.jax.topk_pathsSource ↗

topk_paths(graph: Hypergraph, edge_weights: Array, k: int) -> list[tuple[float, tuple[int, ...]]]

Exact k-best derivations (score, edge-tuple) by score, descending (hard decode).

Wraps the exact topological k-best DP; not differentiable (the argmax set is a hard decode). For a smoothed objective over these paths use topk_free_energy.

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.
kpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnslist[tuple[float, tuple[int, ...]]]

Read the implementation

def topk_paths(
    graph: Hypergraph, edge_weights: Array, k: int
) -> list[tuple[float, tuple[int, ...]]]:
    r"""Exact ``k``-best derivations ``(score, edge-tuple)`` by score, descending (hard decode).

    Wraps the exact topological k-best DP; not differentiable (the argmax set is a hard
    decode). For a smoothed objective over these paths use :func:`topk_free_energy`.
    """
    weights = [float(x) for x in np.asarray(edge_weights).reshape(-1)]
    result: list[tuple[float, tuple[int, ...]]] = kbest_derivations(graph, weights, k)
    return result

Related APIsomnibias.struct.kbest_derivations↗

Other public import paths

• omnibias.struct.jax.topk_paths
• omnibias.struct.jax.distributions.topk_paths

← Package overviewRunnable learning examples →