omnibias.struct.torch.monotonic
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.torch.monotonicModule contract & mathematical conventions
Differentiable Monotonic Alignment Search (torch).
Bit-identical twin of omnibias.struct.jax.monotonic (float64). Soft-MAS replaces the
max of the MAS recursion Q[i, j] = S[i, j] + max(Q[i-1, j-1], Q[i, j-1]) with the
shared lse_beta, unrolled so autograd 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 autograd. 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.
soft_mas
omnibias.struct.torch.soft_masSource ↗soft_mas(score: Tensor, beta: float=1.0) -> TensorSoft-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 | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted regression · soft mas matches oracle and marginals equal autograd
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_mas_matches_oracle_and_marginals_equal_autograd() -> None:
rng = np.random.default_rng(2)
score = rng.standard_normal((3, 5))
st = torch.tensor(score)
hard = hard_mas(score)
assert abs(hard - brute_force_mas(score)) < 1e-9
for beta in (0.5, 2.0, 8.0):
assert abs(float(tmono.soft_mas(st, beta)) - brute_force_soft_mas(score, beta)) < 1e-9
stg = torch.tensor(score, requires_grad=True)
tmono.soft_mas(stg, 3.0).backward()
marg = tmono.soft_mas_marginals(torch.tensor(score), 3.0)
assert torch.max(torch.abs(stg.grad - marg)).item() < 1e-9
assert torch.allclose(marg.sum(dim=0), torch.ones(5), atol=1e-9) # each frame assigned oncepython -m pytest "packages/omnibias-struct/tests/test_plan_attn.py::test_soft_mas_matches_oracle_and_marginals_equal_autograd" -qRegression example · soft mas certified gap and 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_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-9python -m pytest packages/omnibias-struct/tests/test_plan_attn.py -k test_soft_mas_certified_gap_and_parity -qRead the implementation
def soft_mas(score: Tensor, beta: float = 1.0) -> Tensor:
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.torch.soft_masomnibias.struct.torch.monotonic.soft_mas
soft_mas_batched
omnibias.struct.torch.soft_mas_batchedSource ↗soft_mas_batched(score: Tensor, beta: float=1.0) -> TensorBatched soft_mas -> (B,) for score (B, L, T) (via torch.func.vmap).
Parameters
| Name | Type & default | Contract |
|---|---|---|
scorepositional-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 mas and value iteration 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_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() < ATOLpython -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_soft_mas_and_value_iteration_batched_torch" -qRead the implementation
def soft_mas_batched(score: Tensor, beta: float = 1.0) -> Tensor:
r"""Batched :func:`soft_mas` -> ``(B,)`` for ``score`` ``(B, L, T)`` (via ``torch.func.vmap``)."""
from torch.func import vmap
def fwd(s: Tensor) -> Tensor:
return soft_mas(s, beta)
out: Tensor = vmap(fwd)(score)
return outOther public import paths
omnibias.struct.torch.soft_mas_batchedomnibias.struct.torch.monotonic.soft_mas_batched
soft_mas_marginals
omnibias.struct.torch.soft_mas_marginalsSource ↗soft_mas_marginals(score: Tensor, beta: float=1.0) -> TensorClosed-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 | TensorRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
TensorExecuted regression · soft mas matches oracle and marginals equal autograd
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_mas_matches_oracle_and_marginals_equal_autograd() -> None:
rng = np.random.default_rng(2)
score = rng.standard_normal((3, 5))
st = torch.tensor(score)
hard = hard_mas(score)
assert abs(hard - brute_force_mas(score)) < 1e-9
for beta in (0.5, 2.0, 8.0):
assert abs(float(tmono.soft_mas(st, beta)) - brute_force_soft_mas(score, beta)) < 1e-9
stg = torch.tensor(score, requires_grad=True)
tmono.soft_mas(stg, 3.0).backward()
marg = tmono.soft_mas_marginals(torch.tensor(score), 3.0)
assert torch.max(torch.abs(stg.grad - marg)).item() < 1e-9
assert torch.allclose(marg.sum(dim=0), torch.ones(5), atol=1e-9) # each frame assigned oncepython -m pytest "packages/omnibias-struct/tests/test_plan_attn.py::test_soft_mas_matches_oracle_and_marginals_equal_autograd" -qRegression example · structured attention parity and matches marginals
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_structured_attention_parity_and_matches_marginals() -> None:
rng = np.random.default_rng(5)
scores = rng.standard_normal((5, 3))
trans = rng.standard_normal((3, 3))
beta = 1.5
a_t = tattn(torch.tensor(scores), torch.tensor(trans), beta).numpy()
a_j = np.asarray(jattn(jnp.asarray(scores), jnp.asarray(trans), beta))
assert np.max(np.abs(a_t - a_j)) < 1e-9
# It *is* the linear-chain marginal, by construction.
marg = np.asarray(jsvm(jnp.asarray(scores), jnp.asarray(trans), beta))
assert np.max(np.abs(a_j - marg)) < 1e-9
_ = jsoftmax # jax softmax import kept for symmetry with the torch oraclepython -m pytest packages/omnibias-struct/tests/test_plan_attn.py -k test_structured_attention_parity_and_matches_marginals -qRead the implementation
def soft_mas_marginals(score: Tensor, beta: float = 1.0) -> Tensor:
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 = [
torch.stack([torch.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 torch.stack(rows, dim=0)Other public import paths
omnibias.struct.torch.soft_mas_marginalsomnibias.struct.torch.monotonic.soft_mas_marginals