API REFERENCE/omnibias-struct

omnibias.struct.jax.select

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.jax.select

Module contract & mathematical conventions

Certified argmax / measure-mode collapse: soft selection + Gibbs moments (jax).

Bit-identical twin of omnibias.struct.torch.select (float64 -- enable jax_enable_x64). The Gibbs law p_beta(i) ∝ exp(beta s_i) over N logits collapses onto a Dirac at the mode as beta -> inf; certified_argmax pairs the relaxed selection with a sound, closed-form ~omnibias.struct.SelectionCertificate. The moments reuse the closed-form log-sum-exp tower (omnibias.struct.jax._logsumexp): the Gibbs mean is softmax, the covariance is hessian(lse_beta)/beta, and the higher directional cumulants are the exact log-partition jet kappa_m(v_I) = m! * jet_m / beta^m obtained with omnibias.jax.jet.compose_jet -- no autodiff, no finite differences.

**Two axes, never conflated.** The beta -> inf annealing here is the *feasibility / temperature* sense of collapse (the same axis as omnibias-discrete / omnibias-qubo) -- it is **not** the founding delta -> 0 bias collapse. The founding derivative tower is only the exact engine that differentiates lse_beta (softplus^(n) = sigma^(n-1)); 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.

Module source ↗Source revision 9a2d73df0573Running reference examples →

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
logsumexp_beta	omnibias.struct.jax.logsumexp_beta
softmax_beta	omnibias.struct.jax.softmax_beta

FUNCTION

certified_argmax

omnibias.struct.jax.certified_argmaxSource ↗

certified_argmax(logits: Array, beta: float, *, eps: float | None=None) -> tuple[Array, SelectionCertificate]

Relaxed selection softmax(beta * logits) plus its SelectionCertificate.

logits must be 1-D (N choices). Returns (soft_output, certificate) where the certificate bounds the beta -> inf Gibbs-to-Dirac collapse in closed form (value gap log(N)/beta, mode-mass concentration, and -- when eps is given -- L^inf argmax-stability margin > 2 eps). The soft output is differentiable; the certificate is a frozen numpy object suitable for ~omnibias.struct.seal_selection_certificate.

Parameters

Name	Type & default	Contract
logitspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatRequired	See the operation contract and implementation below.
epskeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returnstuple[Array, SelectionCertificate]

Explicit validation & failure cases

• When lg.ndim != 1, raises ValueError(f'certified_argmax expects 1-D logits (N choices), got shape {tuple(lg.shape)}').

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

Executed regression · select torch jax parity

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_select_torch_jax_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import select as js
    from omnibias.struct.torch import select as ts

    for seed in range(6):
        rng = np.random.default_rng(500 + seed)
        # Well-determined instance: a clear margin so the parity tol is not masking a tie.
        s = np.sort(rng.standard_normal(7))[::-1] + np.linspace(0.0, 0.3, 7)[::-1]
        v = rng.standard_normal(7)
        beta = float(rng.uniform(1.0, 12.0))
        st, vt = torch.tensor(s, dtype=torch.float64), torch.tensor(v, dtype=torch.float64)
        sj, vj = jnp.asarray(s), jnp.asarray(v)

        def close(a, b) -> bool:  # noqa: ANN001
            return bool(np.max(np.abs(np.asarray(a) - np.asarray(b))) < PARITY_TOL)

        assert close(ts.soft_max_value(st, beta).numpy(), js.soft_max_value(sj, beta))
        assert close(ts.soft_argmax(st, beta).numpy(), js.soft_argmax(sj, beta))
        assert close(ts.gibbs_mean(st, beta).numpy(), js.gibbs_mean(sj, beta))
        assert close(ts.gibbs_covariance(st, beta).numpy(), js.gibbs_covariance(sj, beta))
        kt = ts.gibbs_cumulants_directional(st, vt, beta, order=4).numpy()
        kj = np.asarray(js.gibbs_cumulants_directional(sj, vj, beta, order=4))
        assert close(kt, kj)
        soft_t, cert_t = ts.certified_argmax(st, beta, eps=0.05)
        soft_j, cert_j = js.certified_argmax(sj, beta, eps=0.05)
        assert close(soft_t.numpy(), soft_j)
        assert cert_t.argmax == cert_j.argmax
        assert cert_t.p_max_lower == pytest.approx(cert_j.p_max_lower, rel=RTOL, abs=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_select.py::test_select_torch_jax_parity" -q

Read the implementation

def certified_argmax(
    logits: Array, beta: float, *, eps: float | None = None
) -> tuple[Array, SelectionCertificate]:
    r"""Relaxed selection ``softmax(beta * logits)`` plus its :class:`SelectionCertificate`.

    ``logits`` must be 1-D (``N`` choices). Returns ``(soft_output, certificate)`` where the
    certificate bounds the ``beta -> inf`` Gibbs-to-Dirac collapse in closed form (value gap
    ``log(N)/beta``, mode-mass concentration, and -- when ``eps`` is given -- ``L^inf``
    argmax-stability ``margin > 2 eps``). The soft output is differentiable; the certificate is
    a frozen numpy object suitable for :func:`~omnibias.struct.seal_selection_certificate`.
    """
    lg = jnp.asarray(logits)
    if lg.ndim != 1:
        raise ValueError(f"certified_argmax expects 1-D logits (N choices), got shape {tuple(lg.shape)}")
    soft = softmax_beta(lg, beta, axis=-1)
    cert = certify_argmax(np.asarray(lg), beta, eps=eps)
    return soft, cert

Related APIsomnibias.struct.certify_argmax↗omnibias.struct.jax.softmax_beta↗

Other public import paths

• omnibias.struct.jax.certified_argmax
• omnibias.struct.jax.select.certified_argmax

FUNCTION

gibbs_covariance

omnibias.struct.jax.gibbs_covarianceSource ↗

gibbs_covariance(logits: Array, beta: float=1.0, *, axis: int=-1) -> Array

The Gibbs one-hot covariance diag(p) - p p^T (= hessian(lse_beta) / beta).

Returns a (..., N, N) array; the exact second moment of the categorical indicator, obtained in closed form from logsumexp_beta_hessian (never autodiff).

Parameters

Name	Type & default	Contract
logitspositional-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 · select torch jax parity

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_select_torch_jax_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import select as js
    from omnibias.struct.torch import select as ts

    for seed in range(6):
        rng = np.random.default_rng(500 + seed)
        # Well-determined instance: a clear margin so the parity tol is not masking a tie.
        s = np.sort(rng.standard_normal(7))[::-1] + np.linspace(0.0, 0.3, 7)[::-1]
        v = rng.standard_normal(7)
        beta = float(rng.uniform(1.0, 12.0))
        st, vt = torch.tensor(s, dtype=torch.float64), torch.tensor(v, dtype=torch.float64)
        sj, vj = jnp.asarray(s), jnp.asarray(v)

        def close(a, b) -> bool:  # noqa: ANN001
            return bool(np.max(np.abs(np.asarray(a) - np.asarray(b))) < PARITY_TOL)

        assert close(ts.soft_max_value(st, beta).numpy(), js.soft_max_value(sj, beta))
        assert close(ts.soft_argmax(st, beta).numpy(), js.soft_argmax(sj, beta))
        assert close(ts.gibbs_mean(st, beta).numpy(), js.gibbs_mean(sj, beta))
        assert close(ts.gibbs_covariance(st, beta).numpy(), js.gibbs_covariance(sj, beta))
        kt = ts.gibbs_cumulants_directional(st, vt, beta, order=4).numpy()
        kj = np.asarray(js.gibbs_cumulants_directional(sj, vj, beta, order=4))
        assert close(kt, kj)
        soft_t, cert_t = ts.certified_argmax(st, beta, eps=0.05)
        soft_j, cert_j = js.certified_argmax(sj, beta, eps=0.05)
        assert close(soft_t.numpy(), soft_j)
        assert cert_t.argmax == cert_j.argmax
        assert cert_t.p_max_lower == pytest.approx(cert_j.p_max_lower, rel=RTOL, abs=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_select.py::test_select_torch_jax_parity" -q

Read the implementation

def gibbs_covariance(logits: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""The Gibbs one-hot covariance ``diag(p) - p p^T`` (``= hessian(lse_beta) / beta``).

    Returns a ``(..., N, N)`` array; the exact second moment of the categorical indicator,
    obtained in closed form from :func:`logsumexp_beta_hessian` (never autodiff).
    """
    return logsumexp_beta_hessian(logits, beta, axis=axis) / beta

Related APIsomnibias.struct.jax.logsumexp_beta_hessian↗

Other public import paths

• omnibias.struct.jax.gibbs_covariance
• omnibias.struct.jax.select.gibbs_covariance

FUNCTION

gibbs_cumulants_directional

omnibias.struct.jax.gibbs_cumulants_directionalSource ↗

gibbs_cumulants_directional(logits: Array, v: Array, beta: float=1.0, order: int=1) -> Array

Exact cumulants kappa_1..kappa_order of X = v_I under the Gibbs law I ~ p.

kappa_1 = E_p[v_I] = <p, v> (the directional mean), kappa_2 = Var_p(v_I) (= v^T (diag(p) - p p^T) v), and higher cumulants follow. They are the derivatives of the log-partition phi(t) = lse_beta(logits + t v): writing phi(t) = lse_beta(logits) + beta^-1 log sum_i p_i e^{t beta v_i}, the jet of the inner log-sum is composed from the analytic log derivative tower via omnibias.jax.jet.compose_jet (Faà di Bruno), giving kappa_m = m! * F_jet[m] / beta^m exactly -- no autodiff. Reduces over the last axis; returns a (order, ...) array (index m-1 holds kappa_m).

Parameters

Name	Type & default	Contract
logitspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
vpositional-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 < 1, raises ValueError(f'order must be >= 1, got {order}').

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

Executed regression · select torch jax parity

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_select_torch_jax_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import select as js
    from omnibias.struct.torch import select as ts

    for seed in range(6):
        rng = np.random.default_rng(500 + seed)
        # Well-determined instance: a clear margin so the parity tol is not masking a tie.
        s = np.sort(rng.standard_normal(7))[::-1] + np.linspace(0.0, 0.3, 7)[::-1]
        v = rng.standard_normal(7)
        beta = float(rng.uniform(1.0, 12.0))
        st, vt = torch.tensor(s, dtype=torch.float64), torch.tensor(v, dtype=torch.float64)
        sj, vj = jnp.asarray(s), jnp.asarray(v)

        def close(a, b) -> bool:  # noqa: ANN001
            return bool(np.max(np.abs(np.asarray(a) - np.asarray(b))) < PARITY_TOL)

        assert close(ts.soft_max_value(st, beta).numpy(), js.soft_max_value(sj, beta))
        assert close(ts.soft_argmax(st, beta).numpy(), js.soft_argmax(sj, beta))
        assert close(ts.gibbs_mean(st, beta).numpy(), js.gibbs_mean(sj, beta))
        assert close(ts.gibbs_covariance(st, beta).numpy(), js.gibbs_covariance(sj, beta))
        kt = ts.gibbs_cumulants_directional(st, vt, beta, order=4).numpy()
        kj = np.asarray(js.gibbs_cumulants_directional(sj, vj, beta, order=4))
        assert close(kt, kj)
        soft_t, cert_t = ts.certified_argmax(st, beta, eps=0.05)
        soft_j, cert_j = js.certified_argmax(sj, beta, eps=0.05)
        assert close(soft_t.numpy(), soft_j)
        assert cert_t.argmax == cert_j.argmax
        assert cert_t.p_max_lower == pytest.approx(cert_j.p_max_lower, rel=RTOL, abs=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_select.py::test_select_torch_jax_parity" -q

Read the implementation

def gibbs_cumulants_directional(
    logits: Array, v: Array, beta: float = 1.0, order: int = 1
) -> Array:
    r"""Exact cumulants ``kappa_1..kappa_order`` of ``X = v_I`` under the Gibbs law ``I ~ p``.

    ``kappa_1 = E_p[v_I] = <p, v>`` (the directional mean), ``kappa_2 = Var_p(v_I)``
    (``= v^T (diag(p) - p p^T) v``), and higher cumulants follow. They are the derivatives of
    the log-partition ``phi(t) = lse_beta(logits + t v)``: writing
    ``phi(t) = lse_beta(logits) + beta^-1 log sum_i p_i e^{t beta v_i}``, the jet of the inner
    ``log-sum`` is composed from the analytic ``log`` derivative tower via
    :func:`omnibias.jax.jet.compose_jet` (Faà di Bruno), giving
    ``kappa_m = m! * F_jet[m] / beta^m`` exactly -- no autodiff. Reduces over the last axis;
    returns a ``(order, ...)`` array (index ``m-1`` holds ``kappa_m``).
    """
    if order < 1:
        raise ValueError(f"order must be >= 1, got {order}")
    s = jnp.asarray(logits)
    vv = jnp.asarray(v)
    p = softmax_beta(s, beta, axis=-1)
    w = beta * vv
    # Inner jet (coefficient form): S_coeffs[k] = sum_i p_i w_i^k / k!  (S(0) = sum_i p_i = 1).
    s_rows = [p.sum(axis=-1)]
    wk = jnp.ones_like(w)
    fact = 1.0
    for k in range(1, order + 1):
        wk = wk * w
        fact *= k
        s_rows.append((p * wk).sum(axis=-1) / fact)
    s_coeffs = jnp.stack(s_rows, axis=0)
    s0 = s_coeffs[0]
    # Analytic log derivative tower at S0: f^(k)(S0) = (-1)^(k-1) (k-1)! / S0^k (k >= 1).
    log_rows = [jnp.log(s0)]
    for k in range(1, order + 1):
        coef = ((-1.0) ** (k - 1)) * math.factorial(k - 1)
        log_rows.append(coef / (s0**k))
    log_tower = jnp.stack(log_rows, axis=0)
    f_jet = compose_jet(s_coeffs, log_tower)
    cumulants = [math.factorial(m) * f_jet[m] / (beta**m) for m in range(1, order + 1)]
    return jnp.stack(cumulants, axis=0)

Related APIsomnibias.jax.compose_jet↗omnibias.struct.jax.softmax_beta↗

Other public import paths

• omnibias.struct.jax.gibbs_cumulants_directional
• omnibias.struct.jax.select.gibbs_cumulants_directional

FUNCTION

gibbs_mean

omnibias.struct.jax.gibbs_meanSource ↗

gibbs_mean(logits: Array, beta: float=1.0, *, axis: int=-1) -> Array

The Gibbs law p_i = softmax(beta * logits)_i -- the mean of the one-hot indicator.

Identical array to soft_argmax; named for the measure/moment view (it is the gradient of soft_max_value, i.e. E_p[e_i]).

Parameters

Name	Type & default	Contract
logitspositional-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 · select torch jax parity

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_select_torch_jax_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import select as js
    from omnibias.struct.torch import select as ts

    for seed in range(6):
        rng = np.random.default_rng(500 + seed)
        # Well-determined instance: a clear margin so the parity tol is not masking a tie.
        s = np.sort(rng.standard_normal(7))[::-1] + np.linspace(0.0, 0.3, 7)[::-1]
        v = rng.standard_normal(7)
        beta = float(rng.uniform(1.0, 12.0))
        st, vt = torch.tensor(s, dtype=torch.float64), torch.tensor(v, dtype=torch.float64)
        sj, vj = jnp.asarray(s), jnp.asarray(v)

        def close(a, b) -> bool:  # noqa: ANN001
            return bool(np.max(np.abs(np.asarray(a) - np.asarray(b))) < PARITY_TOL)

        assert close(ts.soft_max_value(st, beta).numpy(), js.soft_max_value(sj, beta))
        assert close(ts.soft_argmax(st, beta).numpy(), js.soft_argmax(sj, beta))
        assert close(ts.gibbs_mean(st, beta).numpy(), js.gibbs_mean(sj, beta))
        assert close(ts.gibbs_covariance(st, beta).numpy(), js.gibbs_covariance(sj, beta))
        kt = ts.gibbs_cumulants_directional(st, vt, beta, order=4).numpy()
        kj = np.asarray(js.gibbs_cumulants_directional(sj, vj, beta, order=4))
        assert close(kt, kj)
        soft_t, cert_t = ts.certified_argmax(st, beta, eps=0.05)
        soft_j, cert_j = js.certified_argmax(sj, beta, eps=0.05)
        assert close(soft_t.numpy(), soft_j)
        assert cert_t.argmax == cert_j.argmax
        assert cert_t.p_max_lower == pytest.approx(cert_j.p_max_lower, rel=RTOL, abs=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_select.py::test_select_torch_jax_parity" -q

Read the implementation

def gibbs_mean(logits: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""The Gibbs law ``p_i = softmax(beta * logits)_i`` -- the mean of the one-hot indicator.

    Identical array to :func:`soft_argmax`; named for the measure/moment view (it is the
    gradient of :func:`soft_max_value`, i.e. ``E_p[e_i]``).
    """
    return softmax_beta(logits, beta, axis=axis)

Related APIsomnibias.struct.jax.softmax_beta↗

Other public import paths

• omnibias.struct.jax.gibbs_mean
• omnibias.struct.jax.select.gibbs_mean

FUNCTION

soft_argmax

omnibias.struct.jax.soft_argmaxSource ↗

soft_argmax(logits: Array, beta: float=1.0, *, axis: int=-1) -> Array

The relaxed one-hot selection softmax(beta * logits) (-> e_argmax as beta -> inf).

Parameters

Name	Type & default	Contract
logitspositional-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 · select torch jax parity

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_select_torch_jax_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import select as js
    from omnibias.struct.torch import select as ts

    for seed in range(6):
        rng = np.random.default_rng(500 + seed)
        # Well-determined instance: a clear margin so the parity tol is not masking a tie.
        s = np.sort(rng.standard_normal(7))[::-1] + np.linspace(0.0, 0.3, 7)[::-1]
        v = rng.standard_normal(7)
        beta = float(rng.uniform(1.0, 12.0))
        st, vt = torch.tensor(s, dtype=torch.float64), torch.tensor(v, dtype=torch.float64)
        sj, vj = jnp.asarray(s), jnp.asarray(v)

        def close(a, b) -> bool:  # noqa: ANN001
            return bool(np.max(np.abs(np.asarray(a) - np.asarray(b))) < PARITY_TOL)

        assert close(ts.soft_max_value(st, beta).numpy(), js.soft_max_value(sj, beta))
        assert close(ts.soft_argmax(st, beta).numpy(), js.soft_argmax(sj, beta))
        assert close(ts.gibbs_mean(st, beta).numpy(), js.gibbs_mean(sj, beta))
        assert close(ts.gibbs_covariance(st, beta).numpy(), js.gibbs_covariance(sj, beta))
        kt = ts.gibbs_cumulants_directional(st, vt, beta, order=4).numpy()
        kj = np.asarray(js.gibbs_cumulants_directional(sj, vj, beta, order=4))
        assert close(kt, kj)
        soft_t, cert_t = ts.certified_argmax(st, beta, eps=0.05)
        soft_j, cert_j = js.certified_argmax(sj, beta, eps=0.05)
        assert close(soft_t.numpy(), soft_j)
        assert cert_t.argmax == cert_j.argmax
        assert cert_t.p_max_lower == pytest.approx(cert_j.p_max_lower, rel=RTOL, abs=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_select.py::test_select_torch_jax_parity" -q

Read the implementation

def soft_argmax(logits: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""The relaxed one-hot selection ``softmax(beta * logits)`` (``-> e_argmax`` as ``beta -> inf``)."""
    return softmax_beta(logits, beta, axis=axis)

Related APIsomnibias.struct.jax.softmax_beta↗

Other public import paths

• omnibias.struct.jax.soft_argmax
• omnibias.struct.jax.select.soft_argmax

FUNCTION

soft_max_value

omnibias.struct.jax.soft_max_valueSource ↗

soft_max_value(logits: Array, beta: float=1.0, *, axis: int=-1) -> Array

The relaxed max value lse_beta(logits) (>= max, -> max as beta -> inf).

Parameters

Name	Type & default	Contract
logitspositional-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 · select torch jax parity

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_select_torch_jax_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.struct.jax import select as js
    from omnibias.struct.torch import select as ts

    for seed in range(6):
        rng = np.random.default_rng(500 + seed)
        # Well-determined instance: a clear margin so the parity tol is not masking a tie.
        s = np.sort(rng.standard_normal(7))[::-1] + np.linspace(0.0, 0.3, 7)[::-1]
        v = rng.standard_normal(7)
        beta = float(rng.uniform(1.0, 12.0))
        st, vt = torch.tensor(s, dtype=torch.float64), torch.tensor(v, dtype=torch.float64)
        sj, vj = jnp.asarray(s), jnp.asarray(v)

        def close(a, b) -> bool:  # noqa: ANN001
            return bool(np.max(np.abs(np.asarray(a) - np.asarray(b))) < PARITY_TOL)

        assert close(ts.soft_max_value(st, beta).numpy(), js.soft_max_value(sj, beta))
        assert close(ts.soft_argmax(st, beta).numpy(), js.soft_argmax(sj, beta))
        assert close(ts.gibbs_mean(st, beta).numpy(), js.gibbs_mean(sj, beta))
        assert close(ts.gibbs_covariance(st, beta).numpy(), js.gibbs_covariance(sj, beta))
        kt = ts.gibbs_cumulants_directional(st, vt, beta, order=4).numpy()
        kj = np.asarray(js.gibbs_cumulants_directional(sj, vj, beta, order=4))
        assert close(kt, kj)
        soft_t, cert_t = ts.certified_argmax(st, beta, eps=0.05)
        soft_j, cert_j = js.certified_argmax(sj, beta, eps=0.05)
        assert close(soft_t.numpy(), soft_j)
        assert cert_t.argmax == cert_j.argmax
        assert cert_t.p_max_lower == pytest.approx(cert_j.p_max_lower, rel=RTOL, abs=ATOL)

python -m pytest "packages/omnibias-struct/tests/test_select.py::test_select_torch_jax_parity" -q

Read the implementation

def soft_max_value(logits: Array, beta: float = 1.0, *, axis: int = -1) -> Array:
    r"""The relaxed max value ``lse_beta(logits)`` (``>= max``, ``-> max`` as ``beta -> inf``)."""
    return logsumexp_beta(logits, beta, axis=axis)

Related APIsomnibias.struct.jax.logsumexp_beta↗

Other public import paths

• omnibias.struct.jax.soft_max_value
• omnibias.struct.jax.select.soft_max_value

← Package overviewRunnable learning examples →