omnibias.struct.jax.attention
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.jax.attentionModule contract & mathematical conventions
Structured attention as linear-chain soft-DP marginals (jax).
Bit-identical twin of omnibias.struct.torch.attention (float64; needs
jax_enable_x64). 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.jax.soft_viterbi_marginals). 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.jax.structured_attentionSource ↗structured_attention(scores: Array, transitions: Array, beta: float=1.0, *, start: Array | None=None) -> ArrayStructured 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 | ArrayRequired | See the operation contract and implementation below. |
transitionspositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
startkeyword-only | Array | NoneDefault: None | See the operation contract and implementation below. |
ArrayExecuted regression · structured attention parity and matches marginals
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_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::test_structured_attention_parity_and_matches_marginals" -qRead the implementation
def structured_attention(
scores: Array,
transitions: Array,
beta: float = 1.0,
*,
start: Array | None = None,
) -> Array:
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.jax.structured_attentionomnibias.struct.jax.attention.structured_attention
structured_attention_batched
omnibias.struct.jax.structured_attention_batchedSource ↗structured_attention_batched(scores: Array, transitions: Array, beta: float=1.0, *, start: Array | None=None) -> ArrayBatched 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 marginals over the leading
batch axis with jax.vmap -- bit-identical to looping.
Parameters
| Name | Type & default | Contract |
|---|---|---|
scorespositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
transitionspositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
startkeyword-only | Array | NoneDefault: None | See the operation contract and implementation below. |
ArrayRead the implementation
def structured_attention_batched(
scores: Array,
transitions: Array,
beta: float = 1.0,
*,
start: Array | None = None,
) -> Array:
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 marginals over the leading
batch axis with ``jax.vmap`` -- bit-identical to looping.
"""
import jax
t_dim = 0 if transitions.ndim == 3 else None
if start is None:
out: Array = jax.vmap(
lambda sc, tr: structured_attention(sc, tr, beta), in_axes=(0, t_dim)
)(scores, transitions)
return out
s_dim = 0 if start.ndim == 2 else None
out_start: Array = jax.vmap(
lambda sc, tr, st: structured_attention(sc, tr, beta, start=st), in_axes=(0, t_dim, s_dim)
)(scores, transitions, start)
return out_startOther public import paths
omnibias.struct.jax.structured_attention_batchedomnibias.struct.jax.attention.structured_attention_batched