API REFERENCE/omnibias-struct

omnibias.struct.jax.monotonic

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.monotonic

Module contract & mathematical conventions

Differentiable Monotonic Alignment Search (jax).

Bit-identical twin of omnibias.struct.torch.monotonic (float64; needs jax_enable_x64). Soft-MAS replaces the max of Q[i, j] = S[i, j] + max(Q[i-1, j-1], Q[i, j-1]) with the shared lse_beta, unrolled so jax.grad flows through it. As beta -> inf (the temperature axis) it anneals to hard MAS; the delta -> 0 tower gives the closed-form alignment marginals (soft_mas_marginals), pinned equal to jax.grad. Do not conflate the two axes.

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

soft_mas

omnibias.struct.jax.soft_masSource ↗

soft_mas(score: Array, beta: float=1.0) -> Array

Soft-MAS value beta^-1 log sum_alignments exp(beta score) of an (L, T) matrix.

The recursive softmax over monotonic-surjective alignments; -> hard MAS as beta -> inf. Differentiable in score.

Parameters

Name	Type & default	Contract
scorepositional-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 · soft mas certified gap and 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_soft_mas_certified_gap_and_parity() -> None:
    from omnibias.struct import count_alignments  # noqa: PLC0415

    rng = np.random.default_rng(3)
    score = rng.standard_normal((3, 5))
    st, sj = torch.tensor(score), jnp.asarray(score)
    hard = hard_mas(score)
    n = count_alignments(3, 5)
    for beta in (1.0, 4.0, 16.0):
        soft = float(tmono.soft_mas(st, beta))
        cert = certify_soft_dp(hard, soft, n, beta, brute_force_value=brute_force_mas(score))
        assert cert.is_sound
        assert abs(soft - float(jmono.soft_mas(sj, beta))) < 1e-9
    m_t = tmono.soft_mas_marginals(st, 4.0).numpy()
    m_j = np.asarray(jmono.soft_mas_marginals(sj, 4.0))
    assert np.max(np.abs(m_t - m_j)) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_plan_attn.py::test_soft_mas_certified_gap_and_parity" -q

Regression example · soft mas and value iteration 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_mas_and_value_iteration_batched_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import (
        soft_mas,
        soft_mas_batched,
        soft_value_iteration,
        soft_value_iteration_batched,
    )

    rng = np.random.default_rng(1)
    score = torch.tensor(rng.standard_normal((4, 3, 4)))
    b_mas = soft_mas_batched(score, 2.0)
    loop_mas = torch.stack([soft_mas(score[i], 2.0) for i in range(4)])
    assert b_mas.shape == (4,)
    assert torch.max(torch.abs(b_mas - loop_mas)).item() < ATOL

    mdp = _mdp()
    rewards = torch.tensor(rng.standard_normal((4, len(mdp.actions))))
    b_vi = soft_value_iteration_batched(rewards, mdp, 1.5)
    loop_vi = torch.stack([soft_value_iteration(rewards[i], mdp, 1.5) for i in range(4)])
    assert torch.max(torch.abs(b_vi - loop_vi)).item() < ATOL

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

Read the implementation

def soft_mas(score: Array, beta: float = 1.0) -> Array:
    r"""Soft-MAS value ``beta^-1 log sum_alignments exp(beta score)`` of an ``(L, T)`` matrix.

    The recursive softmax over monotonic-surjective alignments; ``-> hard MAS`` as
    ``beta -> inf``. Differentiable in ``score``.
    """
    n_tokens, n_frames = score.shape
    return _forward(score, beta)[n_tokens - 1][n_frames - 1]

Other public import paths

• omnibias.struct.jax.soft_mas
• omnibias.struct.jax.monotonic.soft_mas

FUNCTION

soft_mas_batched

omnibias.struct.jax.soft_mas_batchedSource ↗

soft_mas_batched(score: Array, beta: float=1.0) -> Array

Batched soft_mas -> (B,) for score (B, L, T) (via jax.vmap).

Parameters

Name	Type & default	Contract
scorepositional-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 jax

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_jax() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import (
        soft_align,
        soft_align_batched,
        soft_eisner,
        soft_eisner_batched,
        soft_mas,
        soft_mas_batched,
    )

    rng = np.random.default_rng(4)
    a, b = np.array([0, 1, 2]), np.array([0, 2, 1])
    sub = jnp.asarray(rng.standard_normal((4, 3, 3)))
    gap = jnp.asarray(-0.7)
    b_al = np.asarray(soft_align_batched(a, b, sub, gap, 3.0))
    loop_al = np.asarray([float(soft_align(a, b, sub[i], gap, 3.0)) for i in range(4)])
    assert np.max(np.abs(b_al - loop_al)) < ATOL

    score = jnp.asarray(rng.standard_normal((4, 3, 4)))
    b_mas = np.asarray(soft_mas_batched(score, 2.0))
    loop_mas = np.asarray([float(soft_mas(score[i], 2.0)) for i in range(4)])
    assert np.max(np.abs(b_mas - loop_mas)) < ATOL

    arc = jnp.asarray(rng.standard_normal((4, 4, 4)))
    b_ei = np.asarray(soft_eisner_batched(arc, 2.0))
    loop_ei = np.asarray([float(soft_eisner(arc[i], 2.0)) for i in range(4)])
    assert np.max(np.abs(b_ei - loop_ei)) < ATOL

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

Regression example · soft mas and value iteration 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_mas_and_value_iteration_batched_torch() -> None:
    torch = pytest.importorskip("torch")
    torch.set_default_dtype(torch.float64)
    from omnibias.struct.torch import (
        soft_mas,
        soft_mas_batched,
        soft_value_iteration,
        soft_value_iteration_batched,
    )

    rng = np.random.default_rng(1)
    score = torch.tensor(rng.standard_normal((4, 3, 4)))
    b_mas = soft_mas_batched(score, 2.0)
    loop_mas = torch.stack([soft_mas(score[i], 2.0) for i in range(4)])
    assert b_mas.shape == (4,)
    assert torch.max(torch.abs(b_mas - loop_mas)).item() < ATOL

    mdp = _mdp()
    rewards = torch.tensor(rng.standard_normal((4, len(mdp.actions))))
    b_vi = soft_value_iteration_batched(rewards, mdp, 1.5)
    loop_vi = torch.stack([soft_value_iteration(rewards[i], mdp, 1.5) for i in range(4)])
    assert torch.max(torch.abs(b_vi - loop_vi)).item() < ATOL

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

Read the implementation

def soft_mas_batched(score: Array, beta: float = 1.0) -> Array:
    r"""Batched :func:`soft_mas` -> ``(B,)`` for ``score`` ``(B, L, T)`` (via ``jax.vmap``)."""
    import jax

    out: Array = jax.vmap(lambda s: soft_mas(s, beta))(score)
    return out

Related APIsomnibias.struct.jax.soft_mas↗

Other public import paths

• omnibias.struct.jax.soft_mas_batched
• omnibias.struct.jax.monotonic.soft_mas_batched

FUNCTION

soft_mas_marginals

omnibias.struct.jax.soft_mas_marginalsSource ↗

soft_mas_marginals(score: Array, beta: float=1.0) -> Array

Closed-form soft-alignment matrix A[i, j] = P_beta(frame j assigned to token i).

Forward-backward over the MAS lattice with the tower softmax; equals d soft_mas / d score and each frame column sums to 1. Concentrates on the hard MAS alignment as beta -> inf.

Parameters

Name	Type & default	Contract
scorepositional-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 · soft mas certified gap and 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_soft_mas_certified_gap_and_parity() -> None:
    from omnibias.struct import count_alignments  # noqa: PLC0415

    rng = np.random.default_rng(3)
    score = rng.standard_normal((3, 5))
    st, sj = torch.tensor(score), jnp.asarray(score)
    hard = hard_mas(score)
    n = count_alignments(3, 5)
    for beta in (1.0, 4.0, 16.0):
        soft = float(tmono.soft_mas(st, beta))
        cert = certify_soft_dp(hard, soft, n, beta, brute_force_value=brute_force_mas(score))
        assert cert.is_sound
        assert abs(soft - float(jmono.soft_mas(sj, beta))) < 1e-9
    m_t = tmono.soft_mas_marginals(st, 4.0).numpy()
    m_j = np.asarray(jmono.soft_mas_marginals(sj, 4.0))
    assert np.max(np.abs(m_t - m_j)) < 1e-9

python -m pytest "packages/omnibias-struct/tests/test_plan_attn.py::test_soft_mas_certified_gap_and_parity" -q

Read the implementation

def soft_mas_marginals(score: Array, beta: float = 1.0) -> Array:
    r"""Closed-form soft-alignment matrix ``A[i, j] = P_beta(frame j assigned to token i)``.

    Forward-backward over the MAS lattice with the tower softmax; equals
    ``d soft_mas / d score`` and each frame column sums to ``1``. Concentrates on the hard
    MAS alignment as ``beta -> inf``.
    """
    n_tokens, n_frames = score.shape
    fwd = _forward(score, beta)
    bwd = _backward(score, beta)
    value = fwd[n_tokens - 1][n_frames - 1]
    rows = [
        jnp.stack([jnp.exp(beta * (fwd[i][j] + bwd[i][j] - score[i, j] - value)) for j in range(n_frames)])
        for i in range(n_tokens)
    ]
    return jnp.stack(rows, axis=0)

Other public import paths

• omnibias.struct.jax.soft_mas_marginals
• omnibias.struct.jax.monotonic.soft_mas_marginals

← Package overviewRunnable learning examples →