omnibias.struct.torch.attention
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.torch.attentionModule contract & mathematical conventions
Structured attention as linear-chain soft-DP marginals (torch).
Bit-identical twin of omnibias.struct.jax.attention (float64). Ordinary attention is
an independent per-row softmax; **structured attention** couples the rows with a
transition score and reads the attention weights off the linear-chain *marginals*
(omnibias.struct.torch.soft_viterbi_marginals) -- the closed-form forward-backward
softmax of the delta -> 0 tower. With zero transitions it collapses back to plain
row-wise softmax(beta * scores); as beta -> inf it concentrates on the Viterbi path.
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.
structured_attention
omnibias.struct.torch.structured_attentionSource ↗structured_attention(scores: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> TensorStructured attention weights (T, S) -- linear-chain marginals of scores.
scores are per-position, per-key logits (T, S) and transitions (S, S)
couple adjacent positions; the returned weights are P_beta(key s at position t), each
row summing to 1. Zero transitions recovers independent softmax(beta * scores)
attention; a non-zero transition biases attention towards structurally coherent paths.
Parameters
| Name | Type & default | Contract |
|---|---|---|
scorespositional-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. |
TensorExecuted regression · structured attention collapses to softmax without transitions
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_structured_attention_collapses_to_softmax_without_transitions() -> None:
rng = np.random.default_rng(4)
scores = rng.standard_normal((5, 4))
zero = np.zeros((4, 4))
beta = 2.0
attn = tattn(torch.tensor(scores), torch.tensor(zero), beta)
plain = tsoftmax(torch.tensor(scores), beta, axis=1)
assert torch.max(torch.abs(attn - plain)).item() < 1e-9 # zero transitions -> plain softmax
assert torch.allclose(attn.sum(dim=1), torch.ones(5), atol=1e-9)python -m pytest "packages/omnibias-struct/tests/test_plan_attn.py::test_structured_attention_collapses_to_softmax_without_transitions" -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 structured_attention(
scores: Tensor,
transitions: Tensor,
beta: float = 1.0,
*,
start: Tensor | None = None,
) -> Tensor:
r"""Structured attention weights ``(T, S)`` -- linear-chain marginals of ``scores``.
``scores`` are per-position, per-key logits ``(T, S)`` and ``transitions`` ``(S, S)``
couple adjacent positions; the returned weights are ``P_beta(key s at position t)``, each
row summing to ``1``. Zero ``transitions`` recovers independent ``softmax(beta * scores)``
attention; a non-zero transition biases attention towards structurally coherent paths.
"""
return soft_viterbi_marginals(scores, transitions, beta, start=start)Other public import paths
omnibias.struct.torch.structured_attentionomnibias.struct.torch.attention.structured_attention
structured_attention_batched
omnibias.struct.torch.structured_attention_batchedSource ↗structured_attention_batched(scores: Tensor, transitions: Tensor, beta: float=1.0, *, start: Tensor | None=None) -> TensorBatched structured_attention -> (B, T, S) for scores (B, T, S).
transitions may be shared (S, S) or per-example (B, S, S); start may be
None, shared (S,), or per-example (B, S). Maps the linear-chain marginals over
the leading batch axis with torch.func.vmap -- bit-identical to looping.
Parameters
| Name | Type & default | Contract |
|---|---|---|
scorespositional-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. |
TensorExecuted regression · structured attention 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_structured_attention_batched_torch() -> None:
torch = pytest.importorskip("torch")
torch.set_default_dtype(torch.float64)
from omnibias.struct.torch import structured_attention, structured_attention_batched
rng = np.random.default_rng(2)
scores = torch.tensor(rng.standard_normal((4, 5, 3)))
trans = torch.tensor(rng.standard_normal((3, 3)))
batched = structured_attention_batched(scores, trans, 2.0)
loop = torch.stack([structured_attention(scores[i], trans, 2.0) for i in range(4)])
assert batched.shape == (4, 5, 3)
assert torch.max(torch.abs(batched - loop)).item() < ATOL
# per-example transitions
trans_b = torch.tensor(rng.standard_normal((4, 3, 3)))
batched2 = structured_attention_batched(scores, trans_b, 2.0)
loop2 = torch.stack([structured_attention(scores[i], trans_b[i], 2.0) for i in range(4)])
assert torch.max(torch.abs(batched2 - loop2)).item() < ATOLpython -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_structured_attention_batched_torch" -qRegression example · soft align batched matches loop 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.
def test_soft_align_batched_matches_loop_torch() -> None:
torch = pytest.importorskip("torch")
torch.set_default_dtype(torch.float64)
from omnibias.struct.torch import soft_align, soft_align_batched
rng = np.random.default_rng(0)
a, b = np.array([0, 1, 2, 1]), np.array([0, 2, 1])
sub = torch.tensor(rng.standard_normal((5, 3, 3)))
gap = torch.tensor(rng.standard_normal(5)) # per-example
batched = soft_align_batched(a, b, sub, gap, 3.0)
loop = torch.stack([soft_align(a, b, sub[i], gap[i], 3.0) for i in range(5)])
assert batched.shape == (5,)
assert torch.max(torch.abs(batched - loop)).item() < ATOL
# shared scalar gap also works
shared = soft_align_batched(a, b, sub, torch.tensor(-1.0), 3.0)
loop2 = torch.stack([soft_align(a, b, sub[i], torch.tensor(-1.0), 3.0) for i in range(5)])
assert torch.max(torch.abs(shared - loop2)).item() < ATOLpython -m pytest packages/omnibias-struct/tests/test_batch_families.py -k test_soft_align_batched_matches_loop_torch -qRead the implementation
def structured_attention_batched(
scores: Tensor,
transitions: Tensor,
beta: float = 1.0,
*,
start: Tensor | None = None,
) -> Tensor:
r"""Batched :func:`structured_attention` -> ``(B, T, S)`` for ``scores`` ``(B, T, S)``.
``transitions`` may be shared ``(S, S)`` or per-example ``(B, S, S)``; ``start`` may be
``None``, shared ``(S,)``, or per-example ``(B, S)``. Maps the linear-chain marginals over
the leading batch axis with ``torch.func.vmap`` -- bit-identical to looping.
"""
from torch.func import vmap
t_dim = 0 if transitions.dim() == 3 else None
if start is None:
def fwd(sc: Tensor, tr: Tensor) -> Tensor:
return structured_attention(sc, tr, beta)
out: Tensor = vmap(fwd, in_dims=(0, t_dim))(scores, transitions)
return out
s_dim = 0 if start.dim() == 2 else None
def fwd_start(sc: Tensor, tr: Tensor, st: Tensor) -> Tensor:
return structured_attention(sc, tr, beta, start=st)
out_start: Tensor = vmap(fwd_start, in_dims=(0, t_dim, s_dim))(scores, transitions, start)
return out_startOther public import paths
omnibias.struct.torch.structured_attention_batchedomnibias.struct.torch.attention.structured_attention_batched