omnibias.struct.jax.mtt
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.jax.mttModule contract & mathematical conventions
Differentiable non-projective matrix-tree dependency parsing (JAX).
Bit-identical twin of omnibias.struct.torch.mtt (float64). By Tutte's directed
Matrix-Tree Theorem the partition over the exponentially many spanning arborescences of the
arc graph is a single determinant, so soft_matrix_tree is an **exact** soft value
log det L(beta) / beta (no lse_beta relaxation of a max -- the determinant sums the
trees exactly) and matrix_tree_marginals reads the closed-form arc marginals off
L^{-1} (Koo et al. 2007), pinned equal to autograd. beta -> inf is still the
temperature axis (-> the Chu-Liu/Edmonds maximum arborescence); the difference from the
rest of the package is that here it is exact at finite beta. Column-max stabilised so the
determinant / inverse stay finite at large beta.
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.
matrix_tree_marginals
omnibias.struct.jax.matrix_tree_marginalsSource ↗matrix_tree_marginals(arc: Array, beta: float=1.0) -> ArrayClosed-form arc marginals P_beta(h -> m) via L^{-1} as an (n + 1, n + 1) matrix.
Equals d soft_matrix_tree / d arc (the exact gradient); each modifier column m >= 1
sums to 1 (every word takes exactly one head). As beta -> inf it concentrates on
the maximum arborescence's arcs.
Parameters
| Name | Type & default | Contract |
|---|---|---|
arcpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExecuted regression · mtt marginals match autodiff
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_mtt_marginals_match_autodiff(backend: str) -> None:
pytest.importorskip(backend)
assert probe_mtt_marginals(0, BETA, backend).okpython -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_mtt_marginals_match_autodiff[jax]" -qRegression example · torch jax 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_torch_jax_parity(seed: int) -> None:
arc = _arc(seed, 4)
at, aj = torch.tensor(arc), jnp.asarray(arc)
for beta in (1.0, 8.0):
assert abs(float(st.soft_matrix_tree(at, beta)) - float(sj.soft_matrix_tree(aj, beta))) < 1e-10
mt = st.matrix_tree_marginals(at, beta).detach().numpy()
mj = np.asarray(sj.matrix_tree_marginals(aj, beta))
assert float(np.max(np.abs(mt - mj))) < 1e-10python -m pytest packages/omnibias-struct/tests/test_mtt.py -k test_torch_jax_parity -qRead the implementation
def matrix_tree_marginals(arc: Array, beta: float = 1.0) -> Array:
r"""Closed-form arc marginals ``P_beta(h -> m)`` via ``L^{-1}`` as an ``(n + 1, n + 1)`` matrix.
Equals ``d soft_matrix_tree / d arc`` (the exact gradient); each modifier column ``m >= 1``
sums to ``1`` (every word takes exactly one head). As ``beta -> inf`` it concentrates on
the maximum arborescence's arcs.
"""
ltilde, _c, tw = _laplacian_tilde(arc, beta)
binv = jnp.linalg.inv(ltilde)
diag_b = jnp.diagonal(binv)
root_tilde = tw[0, :]
aword = tw[1:, :]
p_root = root_tilde * diag_b # (n,)
p_word = aword * (diag_b[jnp.newaxis, :] - binv.T) # (n, n)
out = jnp.zeros_like(arc)
out = out.at[0, 1:].set(p_root)
out = out.at[1:, 1:].set(p_word)
return outOther public import paths
omnibias.struct.jax.matrix_tree_marginalsomnibias.struct.jax.mtt.matrix_tree_marginals
soft_matrix_tree
omnibias.struct.jax.soft_matrix_treeSource ↗soft_matrix_tree(arc: Array, beta: float=1.0) -> ArrayExact non-projective soft value log det L(beta) / beta of the arc-score matrix.
arc is the (n + 1, n + 1) arc-score matrix (arc[h, m] = head h -> modifier
m; row/column 0 is the ROOT wall). Differentiable in arc; the determinant is
the *exact* sum over all spanning arborescences, and -> the maximum-arborescence
score as beta -> inf.
Parameters
| Name | Type & default | Contract |
|---|---|---|
arcpositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExecuted regression · mtt marginals match autodiff
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_mtt_marginals_match_autodiff(backend: str) -> None:
pytest.importorskip(backend)
assert probe_mtt_marginals(0, BETA, backend).okpython -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_mtt_marginals_match_autodiff[jax]" -qRegression example · torch jax 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_torch_jax_parity(seed: int) -> None:
arc = _arc(seed, 4)
at, aj = torch.tensor(arc), jnp.asarray(arc)
for beta in (1.0, 8.0):
assert abs(float(st.soft_matrix_tree(at, beta)) - float(sj.soft_matrix_tree(aj, beta))) < 1e-10
mt = st.matrix_tree_marginals(at, beta).detach().numpy()
mj = np.asarray(sj.matrix_tree_marginals(aj, beta))
assert float(np.max(np.abs(mt - mj))) < 1e-10python -m pytest packages/omnibias-struct/tests/test_mtt.py -k test_torch_jax_parity -qRead the implementation
def soft_matrix_tree(arc: Array, beta: float = 1.0) -> Array:
r"""Exact non-projective soft value ``log det L(beta) / beta`` of the arc-score matrix.
``arc`` is the ``(n + 1, n + 1)`` arc-score matrix (``arc[h, m]`` = head ``h`` -> modifier
``m``; row/column ``0`` is the ROOT wall). Differentiable in ``arc``; the determinant is
the *exact* sum over all spanning arborescences, and ``-> `` the maximum-arborescence
score as ``beta -> inf``.
"""
ltilde, c, _ = _laplacian_tilde(arc, beta)
_sign, logabsdet = jnp.linalg.slogdet(ltilde)
value: Array = (logabsdet + c.sum()) / beta
return valueOther public import paths
omnibias.struct.jax.soft_matrix_treeomnibias.struct.jax.mtt.soft_matrix_tree