API REFERENCE/omnibias-struct

omnibias.struct.jax.eisner

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.eisner

Module contract & mathematical conventions

Differentiable Eisner projective dependency parsing on the semiring driver (jax).

Bit-identical twin of omnibias.struct.torch.eisner (float64 -- enable jax_enable_x64). soft_eisner is the lse_beta partition over projective trees and eisner_marginals reads off the closed-form arc marginals P_beta(arc h -> m) (the exact gradient of soft_eisner w.r.t. the arc-score matrix, 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

eisner_marginals

omnibias.struct.jax.eisner_marginalsSource ↗

eisner_marginals(arc: Array, beta: float=1.0) -> Array

Closed-form arc marginals P_beta(arc h -> m) as an (n + 1, n + 1) matrix.

Equals d soft_eisner / d arc (the exact gradient); each modifier column m >= 1 sums to 1. Inside-outside via the tower softmax; equal to autograd.

Parameters

Name	Type & default	Contract
arcpositional-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 · eisner 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_eisner_marginals_match_autodiff(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_eisner_marginals(0, BETA, backend).ok

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_eisner_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:
    arc = _arc(seed, 4)
    at, aj = torch.tensor(arc), jnp.asarray(arc)
    for beta in (1.0, 8.0):
        assert abs(float(st.soft_eisner(at, beta)) - float(sj.soft_eisner(aj, beta))) < 1e-11
        mt = st.eisner_marginals(at, beta).detach().numpy()
        mj = np.asarray(sj.eisner_marginals(aj, beta))
        assert float(np.max(np.abs(mt - mj))) < 1e-11

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

Read the implementation

def eisner_marginals(arc: Array, beta: float = 1.0) -> Array:
    r"""Closed-form arc marginals ``P_beta(arc h -> m)`` as an ``(n + 1, n + 1)`` matrix.

    Equals ``d soft_eisner / d arc`` (the exact gradient); each modifier column ``m >= 1``
    sums to ``1``. Inside-outside via the tower softmax; equal to ``autograd``.
    """
    spec = eisner_hypergraph(int(arc.shape[0]) - 1)
    mu = semiring_marginals(spec.graph, _edge_weights(spec, arc), beta)
    out = jnp.zeros_like(arc)
    for e, s in enumerate(spec.edge_specs):
        if s[0] == "arc":
            out = out.at[int(s[1]), int(s[2])].add(mu[e])
    return out

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

Other public import paths

• omnibias.struct.jax.eisner_marginals
• omnibias.struct.jax.eisner.eisner_marginals

FUNCTION

soft_eisner

omnibias.struct.jax.soft_eisnerSource ↗

soft_eisner(arc: Array, beta: float=1.0) -> Array

Soft projective-parse partition beta^-1 log sum_trees exp(beta score) of arc.

arc is the (n + 1, n + 1) arc-score matrix (arc[h, m] = head h -> modifier m; index 0 is the ROOT). Differentiable in arc; -> best projective tree score as beta -> inf.

Parameters

Name	Type & default	Contract
arcpositional-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 · eisner 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_eisner_marginals_match_autodiff(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_eisner_marginals(0, BETA, backend).ok

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

Regression example · soft eisner 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_eisner_matches_bruteforce(seed: int) -> None:
    arc = _arc(seed, 4)
    at, aj = torch.tensor(arc), jnp.asarray(arc)
    for beta in (1.0, 8.0, 64.0):
        ref = brute_force_projective(arc, beta)
        assert abs(float(st.soft_eisner(at, beta)) - ref) < 1e-11
        assert abs(float(sj.soft_eisner(aj, beta)) - ref) < 1e-11

python -m pytest packages/omnibias-struct/tests/test_eisner.py -k test_soft_eisner_matches_bruteforce -q

Read the implementation

def soft_eisner(arc: Array, beta: float = 1.0) -> Array:
    r"""Soft projective-parse partition ``beta^-1 log sum_trees exp(beta score)`` of ``arc``.

    ``arc`` is the ``(n + 1, n + 1)`` arc-score matrix (``arc[h, m]`` = head ``h`` -> modifier
    ``m``; index ``0`` is the ``ROOT``). Differentiable in ``arc``; ``-> best projective tree
    score`` as ``beta -> inf``.
    """
    spec = eisner_hypergraph(int(arc.shape[0]) - 1)
    return semiring_value(spec.graph, _edge_weights(spec, arc), beta)

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

Other public import paths

• omnibias.struct.jax.soft_eisner
• omnibias.struct.jax.eisner.soft_eisner

FUNCTION

soft_eisner_batched

omnibias.struct.jax.soft_eisner_batchedSource ↗

soft_eisner_batched(arc: Array, beta: float=1.0) -> Array

Batched soft_eisner -> (B,) for arc (B, n + 1, n + 1) (via jax.vmap).

Parameters

Name	Type & default	Contract
arcpositional-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 · families batched 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_families_batched_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 import soft_eisner_batched as jax_ei
    from omnibias.struct.jax import soft_value_iteration_batched as jax_vi
    from omnibias.struct.torch import soft_eisner_batched as torch_ei
    from omnibias.struct.torch import soft_value_iteration_batched as torch_vi

    torch.set_default_dtype(torch.float64)
    rng = np.random.default_rng(5)
    arc = rng.standard_normal((4, 4, 4))
    et = torch_ei(torch.tensor(arc), 2.0).numpy()
    ej = np.asarray(jax_ei(jnp.asarray(arc), 2.0))
    assert np.max(np.abs(et - ej)) < ATOL

    mdp = _mdp()
    rewards = rng.standard_normal((4, len(mdp.actions)))
    vt = torch_vi(torch.tensor(rewards), mdp, 1.5).numpy()
    vj = np.asarray(jax_vi(jnp.asarray(rewards), mdp, 1.5))
    assert np.max(np.abs(vt - vj)) < ATOL

python -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_families_batched_torch_jax_parity" -q

Regression example · soft inside and eisner batched 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_soft_inside_and_eisner_batched_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import (
        soft_eisner,
        soft_eisner_batched,
        soft_inside,
        soft_inside_batched,
    )

    rng = np.random.default_rng(3)
    grammar = _grammar()
    emit = torch.tensor(rng.standard_normal((4, 3, 2)))
    rule = torch.tensor(rng.standard_normal((4, grammar.num_rules)))
    b_in = soft_inside_batched(grammar, emit, rule, 2.0)
    loop_in = torch.stack([soft_inside(grammar, emit[i], rule[i], 2.0) for i in range(4)])
    assert b_in.shape == (4,)
    assert torch.max(torch.abs(b_in - loop_in)).item() < ATOL

    arc = torch.tensor(rng.standard_normal((4, 4, 4)))
    b_ei = soft_eisner_batched(arc, 2.0)
    loop_ei = torch.stack([soft_eisner(arc[i], 2.0) for i in range(4)])
    assert torch.max(torch.abs(b_ei - loop_ei)).item() < ATOL

python -m pytest packages/omnibias-struct/tests/test_batch_families.py -k test_soft_inside_and_eisner_batched_torch -q

Read the implementation

def soft_eisner_batched(arc: Array, beta: float = 1.0) -> Array:
    r"""Batched :func:`soft_eisner` -> ``(B,)`` for ``arc`` ``(B, n + 1, n + 1)`` (via ``jax.vmap``)."""
    import jax

    out: Array = jax.vmap(lambda a: soft_eisner(a, beta))(arc)
    return out

Related APIsomnibias.struct.jax.soft_eisner↗

Other public import paths

• omnibias.struct.jax.soft_eisner_batched
• omnibias.struct.jax.eisner.soft_eisner_batched

← Package overviewRunnable learning examples →