omnibias.struct.torch.eisner
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.torch.eisnerModule contract & mathematical conventions
Differentiable Eisner projective dependency parsing on the semiring driver (torch).
Bit-identical twin of omnibias.struct.jax.eisner (float64). Lifts Eisner's
complete/incomplete span DP onto the shared hypergraph driver
(omnibias.struct.torch.semiring_value / semiring_marginals): soft_eisner
is the lse_beta partition over all projective trees (-> best projective parse as
beta -> inf), 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. The two axes stay apart: beta -> inf is the relaxation,
delta -> 0 differentiates it exactly.
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.
eisner_marginals
omnibias.struct.torch.eisner_marginalsSource ↗eisner_marginals(arc: Tensor, beta: float=1.0) -> TensorClosed-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 (every word has exactly one head). Inside-outside via the tower softmax;
equal to autograd of soft_eisner. As beta -> inf it concentrates on the
best projective tree's arcs.
Parameters
| Name | Type & default | Contract |
|---|---|---|
arcpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted 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.
def test_eisner_marginals_match_autodiff(backend: str) -> None:
pytest.importorskip(backend)
assert probe_eisner_marginals(0, BETA, backend).okpython -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_eisner_marginals_match_autodiff[torch]" -qRegression example · eisner marginals equals autograd
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.
def test_eisner_marginals_equals_autograd(seed: int) -> None:
arc = _arc(seed, 4)
at = torch.tensor(arc, requires_grad=True)
beta = 6.0
value = st.soft_eisner(at, beta)
mu = st.eisner_marginals(at, beta)
(g,) = torch.autograd.grad(value, at)
assert float((mu.detach() - g).abs().max()) < 1e-9python -m pytest packages/omnibias-struct/tests/test_eisner.py -k test_eisner_marginals_equals_autograd -qRead the implementation
def eisner_marginals(arc: Tensor, beta: float = 1.0) -> Tensor:
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`` (every word has exactly one head). Inside-outside via the tower softmax;
equal to ``autograd`` of :func:`soft_eisner`. As ``beta -> inf`` it concentrates on the
best projective tree's arcs.
"""
spec = eisner_hypergraph(int(arc.shape[0]) - 1)
mu = semiring_marginals(spec.graph, _edge_weights(spec, arc), beta)
out = torch.zeros_like(arc)
for e, s in enumerate(spec.edge_specs):
if s[0] == "arc":
out[int(s[1]), int(s[2])] = out[int(s[1]), int(s[2])] + mu[e]
return outOther public import paths
omnibias.struct.torch.eisner_marginalsomnibias.struct.torch.eisner.eisner_marginals
soft_eisner
omnibias.struct.torch.soft_eisnerSource ↗soft_eisner(arc: Tensor, beta: float=1.0) -> TensorSoft 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; row/column 0 is the ROOT). Differentiable in arc; -> best projective
tree score as beta -> inf. The score of a tree is the sum of its arc scores.
Parameters
| Name | Type & default | Contract |
|---|---|---|
arcpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted 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.
def test_eisner_marginals_match_autodiff(backend: str) -> None:
pytest.importorskip(backend)
assert probe_eisner_marginals(0, BETA, backend).okpython -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_eisner_marginals_match_autodiff[torch]" -qRegression example · eisner marginals equals autograd
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.
def test_eisner_marginals_equals_autograd(seed: int) -> None:
arc = _arc(seed, 4)
at = torch.tensor(arc, requires_grad=True)
beta = 6.0
value = st.soft_eisner(at, beta)
mu = st.eisner_marginals(at, beta)
(g,) = torch.autograd.grad(value, at)
assert float((mu.detach() - g).abs().max()) < 1e-9python -m pytest packages/omnibias-struct/tests/test_eisner.py -k test_eisner_marginals_equals_autograd -qRead the implementation
def soft_eisner(arc: Tensor, beta: float = 1.0) -> Tensor:
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``; row/column ``0`` is the ``ROOT``). Differentiable in ``arc``; ``-> best projective
tree score`` as ``beta -> inf``. The score of a tree is the sum of its arc scores.
"""
spec = eisner_hypergraph(int(arc.shape[0]) - 1)
return semiring_value(spec.graph, _edge_weights(spec, arc), beta)Other public import paths
omnibias.struct.torch.soft_eisneromnibias.struct.torch.eisner.soft_eisner
soft_eisner_batched
omnibias.struct.torch.soft_eisner_batchedSource ↗soft_eisner_batched(arc: Tensor, beta: float=1.0) -> TensorBatched soft_eisner -> (B,) for arc (B, n + 1, n + 1) (via torch.func.vmap).
Parameters
| Name | Type & default | Contract |
|---|---|---|
arcpositional-or-keyword | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted regression · soft inside and eisner batched 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.
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() < ATOLpython -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_soft_inside_and_eisner_batched_torch" -qRegression example · families batched 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.
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)) < ATOLpython -m pytest packages/omnibias-struct/tests/test_batch_families.py -k test_families_batched_torch_jax_parity -qRead the implementation
def soft_eisner_batched(arc: Tensor, beta: float = 1.0) -> Tensor:
r"""Batched :func:`soft_eisner` -> ``(B,)`` for ``arc`` ``(B, n + 1, n + 1)`` (via ``torch.func.vmap``)."""
from torch.func import vmap
def fwd(a: Tensor) -> Tensor:
return soft_eisner(a, beta)
out: Tensor = vmap(fwd)(arc)
return outOther public import paths
omnibias.struct.torch.soft_eisner_batchedomnibias.struct.torch.eisner.soft_eisner_batched