API REFERENCE/omnibias-struct

omnibias.struct.jax

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax

Module contract & mathematical conventions

JAX backend for omnibias-struct: soft-DP layers + closed-form lse/softmax jets.

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 →

alignModule ↗attentionModule ↗distributionsModule ↗dtwModule ↗eisnerModule ↗jetsModule ↗monotonicModule ↗mttModule ↗parseModule ↗planModule ↗selectModule ↗semiringModule ↗soft_dpModule ↗tropicalModule ↗

Imports & aliases

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this module	Canonical reference
certified_argmax	omnibias.struct.jax.certified_argmax
chain_lse_jet	omnibias.struct.jax.chain_lse_jet
cky_lse_jet	omnibias.struct.jax.cky_lse_jet
dag_lse_jet	omnibias.struct.jax.dag_lse_jet
dp_value_jet	omnibias.struct.jax.dp_value_jet
eisner_lse_jet	omnibias.struct.jax.eisner_lse_jet
eisner_marginals	omnibias.struct.jax.eisner_marginals
gibbs_covariance	omnibias.struct.jax.gibbs_covariance
gibbs_cumulants_directional	omnibias.struct.jax.gibbs_cumulants_directional
gibbs_mean	omnibias.struct.jax.gibbs_mean
gumbel_relaxed_sample	omnibias.struct.jax.gumbel_relaxed_sample
hypergraph_lse_jet	omnibias.struct.jax.hypergraph_lse_jet
inside_outside	omnibias.struct.jax.inside_outside
lse2_jet	omnibias.struct.jax.lse2_jet
matrix_tree_marginals	omnibias.struct.jax.matrix_tree_marginals
path_entropy	omnibias.struct.jax.path_entropy
sample_paths	omnibias.struct.jax.sample_paths
semiring_marginals	omnibias.struct.jax.semiring_marginals
semiring_value	omnibias.struct.jax.semiring_value
soft_align	omnibias.struct.jax.soft_align
soft_align_batched	omnibias.struct.jax.soft_align_batched
soft_align_marginals	omnibias.struct.jax.soft_align_marginals
soft_argmax	omnibias.struct.jax.soft_argmax
soft_ctc	omnibias.struct.jax.soft_ctc
soft_ctc_batched	omnibias.struct.jax.soft_ctc_batched
soft_ctc_marginals	omnibias.struct.jax.soft_ctc_marginals
soft_dtw	omnibias.struct.jax.soft_dtw
soft_dtw_batched	omnibias.struct.jax.soft_dtw_batched
soft_dtw_marginals	omnibias.struct.jax.soft_dtw_marginals
soft_eisner	omnibias.struct.jax.soft_eisner
soft_eisner_batched	omnibias.struct.jax.soft_eisner_batched
soft_gotoh	omnibias.struct.jax.soft_gotoh
soft_gotoh_marginals	omnibias.struct.jax.soft_gotoh_marginals
soft_inside	omnibias.struct.jax.soft_inside
soft_inside_batched	omnibias.struct.jax.soft_inside_batched
soft_local_align	omnibias.struct.jax.soft_local_align
soft_local_align_marginals	omnibias.struct.jax.soft_local_align_marginals
soft_mas	omnibias.struct.jax.soft_mas
soft_mas_batched	omnibias.struct.jax.soft_mas_batched
soft_mas_marginals	omnibias.struct.jax.soft_mas_marginals
soft_matrix_tree	omnibias.struct.jax.soft_matrix_tree
soft_max_value	omnibias.struct.jax.soft_max_value
soft_shortest_path	omnibias.struct.jax.soft_shortest_path
soft_shortest_path_batched	omnibias.struct.jax.soft_shortest_path_batched
soft_shortest_path_marginals	omnibias.struct.jax.soft_shortest_path_marginals
soft_value_iteration	omnibias.struct.jax.soft_value_iteration
soft_value_iteration_batched	omnibias.struct.jax.soft_value_iteration_batched
soft_viterbi	omnibias.struct.jax.soft_viterbi
soft_viterbi_batched	omnibias.struct.jax.soft_viterbi_batched
soft_viterbi_marginals	omnibias.struct.jax.soft_viterbi_marginals
soft_viterbi_transition_marginals	omnibias.struct.jax.soft_viterbi_transition_marginals
span_marginals	omnibias.struct.jax.span_marginals
structured_attention	omnibias.struct.jax.structured_attention
structured_attention_batched	omnibias.struct.jax.structured_attention_batched
topk_free_energy	omnibias.struct.jax.topk_free_energy
topk_paths	omnibias.struct.jax.topk_paths

FUNCTION

logsumexp_beta

omnibias.struct.jax.logsumexp_betaSource ↗

logsumexp_beta(a: Array, beta: float=1.0, *, axis: int=-1) -> Array

Stable lse_beta(a) = beta^-1 log sum_i exp(beta a_i) along axis.

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

ReturnsArray

Executed regression · sampler empirical marginals converge

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_sampler_empirical_marginals_converge(backend: str) -> None:
    pytest.importorskip(backend)
    assert probe_sampling(0, BETA, backend).ok

python -m pytest "packages/omnibias-struct/tests/test_refinement_families.py::test_sampler_empirical_marginals_converge[jax]" -q

Regression example · pairwise lse matches lse of pair

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_pairwise_lse_matches_lse_of_pair() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import logsumexp_beta, pairwise_lse

    a = jnp.asarray(np.array([0.3, -1.0, 2.0]))
    b = jnp.asarray(np.array([-0.4, 0.5, 1.5]))
    pair = np.asarray(pairwise_lse(a, b, BETA))
    ref = np.asarray(logsumexp_beta(jnp.stack([a, b], axis=-1), BETA, axis=-1))
    assert np.allclose(pair, ref, rtol=RTOL, atol=ATOL)

python -m pytest packages/omnibias-struct/tests/test_logsumexp.py -k test_pairwise_lse_matches_lse_of_pair -q

Read the implementation

def logsumexp_beta(a: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""Stable ``lse_beta(a) = beta^-1 log sum_i exp(beta a_i)`` along ``axis``."""
    scaled = beta * a
    max_a = jnp.max(scaled, axis=axis, keepdims=True)
    log_sum = jnp.squeeze(max_a, axis=axis) + jnp.log(
        jnp.sum(jnp.exp(scaled - max_a), axis=axis),
    )
    return log_sum / beta

Other public import paths

• omnibias.struct.jax.logsumexp_beta
• omnibias.struct.jax.select.logsumexp_beta

FUNCTION

logsumexp_beta_hessian

omnibias.struct.jax.logsumexp_beta_hessianSource ↗

logsumexp_beta_hessian(a: Array, beta: float=1.0, *, axis: int=-1) -> Array

Closed-form Hessian of lse_beta: beta (diag(p) - p p^T) ((..., n, n)).

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

ReturnsArray

Executed regression · jax primitives vs 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.

Full file, imports & fixtures ↗

def test_jax_primitives_vs_autodiff() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import (
        logsumexp_beta,
        logsumexp_beta_hessian,
        logsumexp_beta_jacobian,
    )

    a = jnp.asarray(A_1D)
    grad = np.asarray(jax.grad(lambda x: logsumexp_beta(x, BETA))(a))
    assert np.allclose(grad, np.asarray(logsumexp_beta_jacobian(a, BETA)), rtol=RTOL, atol=ATOL)
    hess = np.asarray(jax.hessian(lambda x: logsumexp_beta(x, BETA))(a))
    assert np.allclose(hess, np.asarray(logsumexp_beta_hessian(a, BETA)), rtol=RTOL, atol=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_logsumexp.py::test_jax_primitives_vs_autodiff" -q

Regression example · torch primitives vs reference and autodiff

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_torch_primitives_vs_reference_and_autodiff() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import (
        logsumexp_beta,
        logsumexp_beta_hessian,
        logsumexp_beta_jacobian,
        softmax_beta,
    )

    a = torch.tensor(A_1D, requires_grad=True)
    val = logsumexp_beta(a, BETA)
    assert np.isclose(val.item(), _ref_lse(A_1D, BETA), rtol=RTOL, atol=ATOL)
    assert np.allclose(softmax_beta(a.detach(), BETA).numpy(), _ref_softmax(A_1D, BETA), rtol=RTOL, atol=ATOL)
    val.backward()
    jac = logsumexp_beta_jacobian(a.detach(), BETA).numpy()
    assert np.allclose(a.grad.numpy(), jac, rtol=RTOL, atol=ATOL)  # grad lse == softmax
    hess = logsumexp_beta_hessian(a.detach(), BETA).numpy()
    p = _ref_softmax(A_1D, BETA)
    assert np.allclose(hess, BETA * (np.diag(p) - np.outer(p, p)), rtol=RTOL, atol=ATOL)

python -m pytest packages/omnibias-struct/tests/test_logsumexp.py -k test_torch_primitives_vs_reference_and_autodiff -q

Read the implementation

def logsumexp_beta_hessian(a: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""Closed-form Hessian of ``lse_beta``: ``beta (diag(p) - p p^T)`` (``(..., n, n)``)."""
    p = softmax_beta(a, beta, axis=axis)
    outer = p[..., :, None] * p[..., None, :]
    n = p.shape[-1]
    diag = jnp.einsum("...i,ij->...ij", p, jnp.eye(n, dtype=p.dtype))
    return beta * (diag - outer)

Related APIsomnibias.struct.jax.softmax_beta↗

FUNCTION

logsumexp_beta_jacobian

omnibias.struct.jax.logsumexp_beta_jacobianSource ↗

logsumexp_beta_jacobian(a: Array, beta: float=1.0, *, axis: int=-1) -> Array

Closed-form gradient of lse_beta w.r.t. a: softmax(beta a).

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

ReturnsArray

Executed regression · jax primitives vs 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.

Full file, imports & fixtures ↗

def test_jax_primitives_vs_autodiff() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import (
        logsumexp_beta,
        logsumexp_beta_hessian,
        logsumexp_beta_jacobian,
    )

    a = jnp.asarray(A_1D)
    grad = np.asarray(jax.grad(lambda x: logsumexp_beta(x, BETA))(a))
    assert np.allclose(grad, np.asarray(logsumexp_beta_jacobian(a, BETA)), rtol=RTOL, atol=ATOL)
    hess = np.asarray(jax.hessian(lambda x: logsumexp_beta(x, BETA))(a))
    assert np.allclose(hess, np.asarray(logsumexp_beta_hessian(a, BETA)), rtol=RTOL, atol=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_logsumexp.py::test_jax_primitives_vs_autodiff" -q

Regression example · torch primitives vs reference and autodiff

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_torch_primitives_vs_reference_and_autodiff() -> None:
    torch = pytest.importorskip("torch")
    from omnibias.struct.torch import (
        logsumexp_beta,
        logsumexp_beta_hessian,
        logsumexp_beta_jacobian,
        softmax_beta,
    )

    a = torch.tensor(A_1D, requires_grad=True)
    val = logsumexp_beta(a, BETA)
    assert np.isclose(val.item(), _ref_lse(A_1D, BETA), rtol=RTOL, atol=ATOL)
    assert np.allclose(softmax_beta(a.detach(), BETA).numpy(), _ref_softmax(A_1D, BETA), rtol=RTOL, atol=ATOL)
    val.backward()
    jac = logsumexp_beta_jacobian(a.detach(), BETA).numpy()
    assert np.allclose(a.grad.numpy(), jac, rtol=RTOL, atol=ATOL)  # grad lse == softmax
    hess = logsumexp_beta_hessian(a.detach(), BETA).numpy()
    p = _ref_softmax(A_1D, BETA)
    assert np.allclose(hess, BETA * (np.diag(p) - np.outer(p, p)), rtol=RTOL, atol=ATOL)

python -m pytest packages/omnibias-struct/tests/test_logsumexp.py -k test_torch_primitives_vs_reference_and_autodiff -q

Read the implementation

def logsumexp_beta_jacobian(a: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""Closed-form gradient of ``lse_beta`` w.r.t. ``a``: ``softmax(beta a)``."""
    return softmax_beta(a, beta, axis=axis)

Related APIsomnibias.struct.jax.softmax_beta↗

FUNCTION

pairwise_lse

omnibias.struct.jax.pairwise_lseSource ↗

pairwise_lse(a: Array, b: Array, beta: float=1.0) -> Array

Elementwise soft-max combine lse_beta(a, b) = beta^-1 log(e^{beta a} + e^{beta b}).

Equal to a + beta^-1 softplus(beta (b - a)); computed by the stable symmetric log-sum-exp of the stacked pair.

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
bpositional-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 · pairwise lse matches lse of pair

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_pairwise_lse_matches_lse_of_pair() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import logsumexp_beta, pairwise_lse

    a = jnp.asarray(np.array([0.3, -1.0, 2.0]))
    b = jnp.asarray(np.array([-0.4, 0.5, 1.5]))
    pair = np.asarray(pairwise_lse(a, b, BETA))
    ref = np.asarray(logsumexp_beta(jnp.stack([a, b], axis=-1), BETA, axis=-1))
    assert np.allclose(pair, ref, rtol=RTOL, atol=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_logsumexp.py::test_pairwise_lse_matches_lse_of_pair" -q

Regression example · pairwise lse jet from tower matches finite difference

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_pairwise_lse_jet_from_tower_matches_finite_difference() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import pairwise_lse, pairwise_lse_jet

    a0, b0, db = jnp.asarray(0.3), jnp.asarray(-0.4), jnp.asarray(1.0)
    jet = np.asarray(pairwise_lse_jet(a0, b0, db, BETA, order=2))

    def f(t: float) -> float:
        return float(pairwise_lse(a0, b0 + t * db, BETA))

    h = 1e-5
    d1 = (f(h) - f(-h)) / (2 * h)
    d2 = (f(h) - 2 * f(0.0) + f(-h)) / (h * h)
    assert np.isclose(jet[0], f(0.0), rtol=RTOL, atol=ATOL)  # c0 = value
    assert np.isclose(jet[1], d1, atol=1e-6)  # c1 = f'(0)
    assert np.isclose(jet[2], d2 / 2.0, atol=1e-4)  # c2 = f''(0) / 2

python -m pytest packages/omnibias-struct/tests/test_logsumexp.py -k test_pairwise_lse_jet_from_tower_matches_finite_difference -q

Read the implementation

def pairwise_lse(a: Array, b: Array, beta: float = 1.0) -> Array:
    r"""Elementwise soft-max combine ``lse_beta(a, b) = beta^-1 log(e^{beta a} + e^{beta b})``.

    Equal to ``a + beta^-1 softplus(beta (b - a))``; computed by the stable symmetric
    log-sum-exp of the stacked pair.
    """
    stacked = jnp.stack([a, b], axis=-1)
    return logsumexp_beta(stacked, beta, axis=-1)

Related APIsomnibias.struct.jax.logsumexp_beta↗

FUNCTION

pairwise_lse_jet

omnibias.struct.jax.pairwise_lse_jetSource ↗

pairwise_lse_jet(a0: Array, b0: Array, db: Array, beta: float=1.0, order: int=1) -> Array

Taylor jet (order order) of t -> lse_beta(a0, b0 + t db) from the tower.

Writing lse_beta(a0, b0 + t db) = a0 + beta^-1 softplus(u(t)) with u(t) = beta (b0 - a0) + (beta db) t, the jet is omnibias.jax.jet.compose_jet of the linear u-jet with the closed-form softplus derivative tower softplus^(k)(u_0) (softplus^(k) = sigma^(k-1) from omnibias.core) -- exact, no autodiff or finite differences. Returns coefficients c_k = f^(k)(0) / k! along a new leading axis of length order + 1.

Parameters

Name	Type & default	Contract
a0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
b0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
dbpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
orderpositional-or-keyword	intDefault: 1	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

• When order < 0, raises ValueError(f'order must be >= 0, got {order}').
• When fastpath is None, raises NotImplementedError('softplus is missing its closed-form fastpath tower').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · pairwise lse jet from tower matches finite difference

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_pairwise_lse_jet_from_tower_matches_finite_difference() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import pairwise_lse, pairwise_lse_jet

    a0, b0, db = jnp.asarray(0.3), jnp.asarray(-0.4), jnp.asarray(1.0)
    jet = np.asarray(pairwise_lse_jet(a0, b0, db, BETA, order=2))

    def f(t: float) -> float:
        return float(pairwise_lse(a0, b0 + t * db, BETA))

    h = 1e-5
    d1 = (f(h) - f(-h)) / (2 * h)
    d2 = (f(h) - 2 * f(0.0) + f(-h)) / (h * h)
    assert np.isclose(jet[0], f(0.0), rtol=RTOL, atol=ATOL)  # c0 = value
    assert np.isclose(jet[1], d1, atol=1e-6)  # c1 = f'(0)
    assert np.isclose(jet[2], d2 / 2.0, atol=1e-4)  # c2 = f''(0) / 2

python -m pytest "packages/omnibias-struct/tests/test_logsumexp.py::test_pairwise_lse_jet_from_tower_matches_finite_difference" -q

Regression example · primitives and jets 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.

Full file, imports & fixtures ↗

def test_primitives_and_jets_torch_jax_parity() -> None:
    pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    import torch
    from omnibias.struct.jax import logsumexp_beta as lj
    from omnibias.struct.jax import pairwise_lse_jet as jj
    from omnibias.struct.torch import logsumexp_beta as lt
    from omnibias.struct.torch import pairwise_lse_jet as jt

    vt = lt(torch.tensor(A_1D), BETA).item()
    vj = float(lj(jnp.asarray(A_1D), BETA))
    assert abs(vt - vj) < ATOL
    f64 = torch.float64
    tj = np.asarray(
        jt(torch.tensor(0.3, dtype=f64), torch.tensor(-0.4, dtype=f64), torch.tensor(1.0, dtype=f64), BETA, order=3).detach()
    )
    jjj = np.asarray(jj(jnp.asarray(0.3), jnp.asarray(-0.4), jnp.asarray(1.0), BETA, order=3))
    assert np.allclose(tj, jjj, rtol=RTOL, atol=ATOL)

python -m pytest packages/omnibias-struct/tests/test_logsumexp.py -k test_primitives_and_jets_torch_jax_parity -q

Read the implementation

def pairwise_lse_jet(
    a0: Array,
    b0: Array,
    db: Array,
    beta: float = 1.0,
    order: int = 1,
) -> Array:
    r"""Taylor jet (order ``order``) of ``t -> lse_beta(a0, b0 + t db)`` from the tower.

    Writing ``lse_beta(a0, b0 + t db) = a0 + beta^-1 softplus(u(t))`` with
    ``u(t) = beta (b0 - a0) + (beta db) t``, the jet is
    :func:`omnibias.jax.jet.compose_jet` of the linear ``u``-jet with the closed-form
    ``softplus`` derivative tower ``softplus^(k)(u_0)`` (``softplus^(k) = sigma^(k-1)``
    from :mod:`omnibias.core`) -- exact, no autodiff or finite differences. Returns
    coefficients ``c_k = f^(k)(0) / k!`` along a new leading axis of length ``order + 1``.
    """
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    spec = get_activation("softplus")
    fastpath = spec.fastpath
    if fastpath is None:  # defensive: softplus always ships its closed-form tower
        raise NotImplementedError("softplus is missing its closed-form fastpath tower")
    a0 = jnp.asarray(a0)
    b0 = jnp.asarray(b0)
    db = jnp.asarray(db)
    u0 = beta * (b0 - a0)
    zero = jnp.zeros_like(u0)
    u_rows = [u0]
    if order >= 1:
        u_rows.append((beta * db) + zero)
    u_rows.extend(zero for _ in range(order - 1))
    u_jet = jnp.stack(u_rows[: order + 1], axis=0)
    tower_rows = [spec.forward(u0)]
    for k in range(1, order + 1):
        tower_rows.append(fastpath(u0, k))
    sigma_tower = jnp.stack(tower_rows, axis=0)
    sp_jet = compose_jet(u_jet, sigma_tower) / beta
    out0 = (sp_jet[0] + a0)[None, ...]
    return jnp.concatenate([out0, sp_jet[1:]], axis=0)

Related APIsomnibias.jax.get_activation↗omnibias.jax.compose_jet↗

FUNCTION

softmax_beta

omnibias.struct.jax.softmax_betaSource ↗

softmax_beta(a: Array, beta: float=1.0, *, axis: int=-1) -> Array

Stable softmax(beta a) along axis -- the gradient of logsumexp_beta.

Parameters

Name	Type & default	Contract
apositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatDefault: 1.0	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

ReturnsArray

Executed regression · jax primitives vs 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.

Full file, imports & fixtures ↗

def test_jax_primitives_vs_autodiff() -> None:
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import (
        logsumexp_beta,
        logsumexp_beta_hessian,
        logsumexp_beta_jacobian,
    )

    a = jnp.asarray(A_1D)
    grad = np.asarray(jax.grad(lambda x: logsumexp_beta(x, BETA))(a))
    assert np.allclose(grad, np.asarray(logsumexp_beta_jacobian(a, BETA)), rtol=RTOL, atol=ATOL)
    hess = np.asarray(jax.hessian(lambda x: logsumexp_beta(x, BETA))(a))
    assert np.allclose(hess, np.asarray(logsumexp_beta_hessian(a, BETA)), rtol=RTOL, atol=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_logsumexp.py::test_jax_primitives_vs_autodiff" -q

Regression 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.

Full file, imports & fixtures ↗

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 oracle

python -m pytest packages/omnibias-struct/tests/test_plan_attn.py -k test_structured_attention_parity_and_matches_marginals -q

Read the implementation

def softmax_beta(a: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""Stable ``softmax(beta a)`` along ``axis`` -- the gradient of :func:`logsumexp_beta`."""
    scaled = beta * a
    max_a = jnp.max(scaled, axis=axis, keepdims=True)
    exp_a = jnp.exp(scaled - max_a)
    return exp_a / jnp.sum(exp_a, axis=axis, keepdims=True)

Other public import paths

• omnibias.struct.jax.softmax_beta
• omnibias.struct.jax.select.softmax_beta

← Package overviewRunnable learning examples →