API REFERENCE/omnibias-struct

omnibias.struct.torch.curvature

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.torch.curvature

Module contract & mathematical conventions

Second-order curvature *through* a soft-DP value (torch bridge to omnibias-curvature).

The soft Viterbi value V(E) = lse_beta over paths of score(E) is a log-sum-exp of functions **linear** in the emissions E -- hence convex, with a positive-semidefinite Hessian. These helpers expose that curvature by routing the unrolled soft-DP layer through the generic exact machinery in omnibias.curvature.torch (hvp / dense_hessian / hessian_top_eigenvalue). The exact Hessian-vector product agrees to machine precision with the closed-form directional curvature 2 * dp_value_jet[2] from omnibias.struct.torch.jets, giving two independent second-order routes.

Torch only: the JAX curvature surface is one-layer-Riccati, so the second-order bridge is not mirrored there (the delta -> 0 closed-form jet in omnibias.struct.jax.jets is the JAX-side higher-order route).

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_directional_curvature

omnibias.struct.torch.chain_directional_curvatureSource ↗

chain_directional_curvature(emissions: Tensor, transitions: Tensor, direction: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Directional curvature direction^T H direction (equals 2 * chain_lse_jet[2]).

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
directionpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · hvp matches closed form jet curvature

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_hvp_matches_closed_form_jet_curvature() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_directional_curvature, chain_lse_jet

    trellis = random_chain(0)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    d = torch.tensor(np.random.default_rng(1).standard_normal(e.shape))
    beta = 3.0
    curv_hvp = float(chain_directional_curvature(e, trn, d, beta, start=st))
    curv_jet = 2.0 * float(chain_lse_jet(e, trn, d, beta, order=2, start=st)[2])
    assert abs(curv_hvp - curv_jet) < ATOL

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_hvp_matches_closed_form_jet_curvature" -q

Read the implementation

def chain_directional_curvature(
    emissions: Tensor,
    transitions: Tensor,
    direction: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Directional curvature ``direction^T H direction`` (equals ``2 * chain_lse_jet[2]``)."""
    _, hv = chain_value_and_hvp(emissions, transitions, direction, beta, start=start)
    return (hv * direction).sum()

Related APIsomnibias.struct.torch.chain_value_and_hvp↗

Other public import paths

• omnibias.struct.torch.chain_directional_curvature
• omnibias.struct.torch.curvature.chain_directional_curvature

FUNCTION

chain_hessian

omnibias.struct.torch.chain_hessianSource ↗

chain_hessian(emissions: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Dense emission Hessian of the soft Viterbi value, shape (T*S, T*S) (PSD).

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · soft viterbi hessian is psd

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_soft_viterbi_hessian_is_psd() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_hessian

    trellis = random_chain(2)
    h = chain_hessian(
        torch.tensor(trellis.emissions), torch.tensor(trellis.transitions), 5.0,
        start=torch.tensor(trellis.start),
    ).numpy()
    assert np.allclose(h, h.T, atol=1e-10)  # symmetric
    eigs = np.linalg.eigvalsh(h)
    assert eigs.min() > -1e-9  # convex -> PSD

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_soft_viterbi_hessian_is_psd" -q

Regression example · sharpness matches dense eigensolve and grows with beta

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_sharpness_matches_dense_eigensolve_and_grows_with_beta() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_hessian, chain_sharpness

    trellis = random_chain(3)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    top = float(chain_sharpness(e, trn, 4.0, start=st))
    ref = float(np.linalg.eigvalsh(chain_hessian(e, trn, 4.0, start=st).numpy())[-1])
    assert abs(top - ref) < 1e-9
    # sharper relaxation at higher beta
    assert float(chain_sharpness(e, trn, 8.0, start=st)) > float(chain_sharpness(e, trn, 1.0, start=st))

python -m pytest packages/omnibias-struct/tests/test_curvature.py -k test_sharpness_matches_dense_eigensolve_and_grows_with_beta -q

Read the implementation

def chain_hessian(
    emissions: Tensor,
    transitions: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Dense emission Hessian of the soft Viterbi value, shape ``(T*S, T*S)`` (PSD)."""
    leaf = emissions.detach().clone().requires_grad_(True)
    loss = soft_viterbi(leaf, transitions, beta, start=start)
    h: Tensor = dense_hessian(loss, [leaf])
    return h

Related APIsomnibias.curvature.torch.dense_hessian↗omnibias.struct.torch.soft_viterbi↗

Other public import paths

• omnibias.struct.torch.chain_hessian
• omnibias.struct.torch.curvature.chain_hessian

FUNCTION

chain_sharpness

omnibias.struct.torch.chain_sharpnessSource ↗

chain_sharpness(emissions: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> Tensor

Top Hessian eigenvalue (sharpness) of the soft Viterbi loss landscape in E.

Non-negative because soft_viterbi is convex in the emissions; grows with beta (the relaxation sharpens towards the hard optimum).

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

ReturnsTensor

Executed regression · sharpness matches dense eigensolve and grows with beta

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_sharpness_matches_dense_eigensolve_and_grows_with_beta() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_hessian, chain_sharpness

    trellis = random_chain(3)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    top = float(chain_sharpness(e, trn, 4.0, start=st))
    ref = float(np.linalg.eigvalsh(chain_hessian(e, trn, 4.0, start=st).numpy())[-1])
    assert abs(top - ref) < 1e-9
    # sharper relaxation at higher beta
    assert float(chain_sharpness(e, trn, 8.0, start=st)) > float(chain_sharpness(e, trn, 1.0, start=st))

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_sharpness_matches_dense_eigensolve_and_grows_with_beta" -q

Read the implementation

def chain_sharpness(
    emissions: Tensor,
    transitions: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> Tensor:
    r"""Top Hessian eigenvalue (sharpness) of the soft Viterbi loss landscape in ``E``.

    Non-negative because ``soft_viterbi`` is convex in the emissions; grows with ``beta``
    (the relaxation sharpens towards the hard optimum).
    """
    top: Tensor = hessian_top_eigenvalue(chain_hessian(emissions, transitions, beta, start=start))
    return top

Related APIsomnibias.curvature.torch.hessian_top_eigenvalue↗omnibias.struct.torch.chain_hessian↗

Other public import paths

• omnibias.struct.torch.chain_sharpness
• omnibias.struct.torch.curvature.chain_sharpness

FUNCTION

chain_value_and_hvp

omnibias.struct.torch.chain_value_and_hvpSource ↗

chain_value_and_hvp(emissions: Tensor, transitions: Tensor, vector: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> tuple[Tensor, Tensor]

Soft Viterbi value and the exact Hessian-vector product H @ vector in E.

H = d^2 soft_viterbi / d emissions^2; vector has the emission shape (T, S) and the returned Hv matches it. Uses omnibias.curvature.torch.hvp (a double backward), so it is exact, not finite-difference.

Parameters

Name	Type & default	Contract
emissionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
transitionspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
vectorpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
startkeyword-only	Tensor | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[Tensor, Tensor]

Executed regression · value and hvp value matches soft viterbi

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_value_and_hvp_value_matches_soft_viterbi() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import chain_value_and_hvp, soft_viterbi

    trellis = random_chain(1)
    e = torch.tensor(trellis.emissions)
    trn = torch.tensor(trellis.transitions)
    st = torch.tensor(trellis.start)
    v = torch.tensor(np.random.default_rng(0).standard_normal(e.shape))
    value, hv = chain_value_and_hvp(e, trn, v, 4.0, start=st)
    assert abs(float(value) - float(soft_viterbi(e, trn, 4.0, start=st))) < 1e-12
    assert hv.shape == e.shape

python -m pytest "packages/omnibias-struct/tests/test_curvature.py::test_value_and_hvp_value_matches_soft_viterbi" -q

Read the implementation

def chain_value_and_hvp(
    emissions: Tensor,
    transitions: Tensor,
    vector: Tensor,
    beta: float = 1.0,
    *,
    start: Tensor | None = None,
) -> tuple[Tensor, Tensor]:
    r"""Soft Viterbi value and the exact Hessian-vector product ``H @ vector`` in ``E``.

    ``H = d^2 soft_viterbi / d emissions^2``; ``vector`` has the emission shape ``(T, S)``
    and the returned ``Hv`` matches it. Uses :func:`omnibias.curvature.torch.hvp` (a double
    backward), so it is exact, not finite-difference.
    """
    leaf = emissions.detach().clone().requires_grad_(True)
    loss = soft_viterbi(leaf, transitions, beta, start=start)
    hv = hvp(loss, [leaf], [vector])
    return loss.detach(), hv[0]

Related APIsomnibias.curvature.torch.hvp↗omnibias.struct.torch.soft_viterbi↗

Other public import paths

• omnibias.struct.torch.chain_value_and_hvp
• omnibias.struct.torch.curvature.chain_value_and_hvp

← Package overviewRunnable learning examples →