API REFERENCE/omnibias-struct

omnibias.struct.decision.torch

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.struct.decision.torch

Module contract & mathematical conventions

Torch DecisionLayer: an embeddable certified argmax-over-N decision layer.

The trainable face of omnibias.struct.decision. forward is the relaxed decision softmax(beta * scores) (differentiable, batched); hard is the beta -> inf decision argmax(scores); certificate returns the closed-form ~omnibias.struct.SelectionCertificate (value gap log(N)/beta, mode-mass concentration, L^inf argmax-stability radius). Reuses omnibias.struct.torch.certified_argmax and soft_argmax -- no new math.

Terminology: the beta -> inf annealing is the feasibility / temperature sense of "collapse", distinct from the founding delta -> 0 bias collapse (sigma^(K-1)).

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 →

CLASS

DecisionLayer

omnibias.struct.decision.torch.DecisionLayerSource ↗

class DecisionLayer(nn.Module)

An embeddable certified decision (argmax-over-N) layer.

Parameters

beta: Inverse temperature of the relaxed decision (larger -> sharper / closer to argmax). eps: Optional default L^inf perturbation radius for the argmax-stability sub-claim. axis: The axis the decision is taken over (default the last).

Parameters

Name	Type & default	Contract
betapositional-or-keyword	floatDefault: 10.0	Inverse temperature of the relaxed decision (larger -> sharper / closer to argmax).
epskeyword-only	float | NoneDefault: None	Optional default L^inf perturbation radius for the argmax-stability sub-claim.
axiskeyword-only	intDefault: -1	The axis the decision is taken over (default the last).

Executed regression · layer certificate matches functional

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_layer_certificate_matches_functional() -> None:
    layer = DecisionLayer(beta=8.0, eps=0.05)
    scores = torch.tensor([1.0, -0.5, 0.3], dtype=torch.float64)
    soft, cert = layer.certified(scores)
    assert torch.allclose(soft, layer.forward(scores))
    ref = certified_decision(scores.numpy(), beta=8.0, eps=0.05)
    assert cert.argmax == ref.argmax
    assert math.isclose(cert.gap_bound, ref.gap_bound, rel_tol=1e-12)
    assert math.isclose(cert.p_max, ref.p_max, rel_tol=1e-9)

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_layer_certificate_matches_functional" -q

Regression example · soft decision is a distribution and hardens

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_soft_decision_is_a_distribution_and_hardens() -> None:
    layer = DecisionLayer(beta=1.0)
    scores = torch.tensor([0.2, 1.5, 0.9, -0.3], dtype=torch.float64)
    p = layer.forward(scores)
    assert torch.allclose(p.sum(), torch.tensor(1.0, dtype=torch.float64))
    assert torch.all(p >= 0.0)
    # as beta -> inf the relaxed decision approaches the one-hot at the argmax (index 1)
    sharp = DecisionLayer(beta=200.0).forward(scores)
    assert torch.argmax(sharp).item() == 1
    assert sharp[1].item() > 0.999
    assert layer.hard(scores).item() == 1

python -m pytest packages/omnibias-struct/tests/test_decision_layer.py -k test_soft_decision_is_a_distribution_and_hardens -q

Fields & methods

__init__method

__init__(self, beta: float=10.0, *, eps: float | None=None, axis: int=-1) -> None

See the class contract above and the exact method implementation below.

Parameters

Name	Type & default	Contract
betapositional-or-keyword	floatDefault: 10.0	See the operation contract and implementation below.
epskeyword-only	float | NoneDefault: None	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When beta <= 0.0, raises ValueError(f'beta must be > 0, got {beta}').

Executed regression · layer certificate matches functional

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_layer_certificate_matches_functional() -> None:
    layer = DecisionLayer(beta=8.0, eps=0.05)
    scores = torch.tensor([1.0, -0.5, 0.3], dtype=torch.float64)
    soft, cert = layer.certified(scores)
    assert torch.allclose(soft, layer.forward(scores))
    ref = certified_decision(scores.numpy(), beta=8.0, eps=0.05)
    assert cert.argmax == ref.argmax
    assert math.isclose(cert.gap_bound, ref.gap_bound, rel_tol=1e-12)
    assert math.isclose(cert.p_max, ref.p_max, rel_tol=1e-9)

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_layer_certificate_matches_functional" -q

Method implementation

    def __init__(self, beta: float = 10.0, *, eps: float | None = None, axis: int = -1) -> None:
        super().__init__()
        if beta <= 0.0:
            raise ValueError(f"beta must be > 0, got {beta}")
        self.beta = float(beta)
        self.eps = None if eps is None else float(eps)
        self.axis = int(axis)

View source ↗

forwardmethod

forward(self, scores: Tensor) -> Tensor

The relaxed decision softmax(beta * scores) (differentiable).

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · layer certificate matches functional

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_layer_certificate_matches_functional() -> None:
    layer = DecisionLayer(beta=8.0, eps=0.05)
    scores = torch.tensor([1.0, -0.5, 0.3], dtype=torch.float64)
    soft, cert = layer.certified(scores)
    assert torch.allclose(soft, layer.forward(scores))
    ref = certified_decision(scores.numpy(), beta=8.0, eps=0.05)
    assert cert.argmax == ref.argmax
    assert math.isclose(cert.gap_bound, ref.gap_bound, rel_tol=1e-12)
    assert math.isclose(cert.p_max, ref.p_max, rel_tol=1e-9)

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_layer_certificate_matches_functional" -q

Method implementation

    def forward(self, scores: Tensor) -> Tensor:
        r"""The relaxed decision ``softmax(beta * scores)`` (differentiable)."""
        out: Tensor = soft_argmax(scores, self.beta, axis=self.axis)
        return out

View source ↗

hardmethod

hard(self, scores: Tensor) -> Tensor

The beta -> inf hard decision argmax(scores) (indices).

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.

Returns: Tensor

Executed regression · soft decision is a distribution and hardens

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_soft_decision_is_a_distribution_and_hardens() -> None:
    layer = DecisionLayer(beta=1.0)
    scores = torch.tensor([0.2, 1.5, 0.9, -0.3], dtype=torch.float64)
    p = layer.forward(scores)
    assert torch.allclose(p.sum(), torch.tensor(1.0, dtype=torch.float64))
    assert torch.all(p >= 0.0)
    # as beta -> inf the relaxed decision approaches the one-hot at the argmax (index 1)
    sharp = DecisionLayer(beta=200.0).forward(scores)
    assert torch.argmax(sharp).item() == 1
    assert sharp[1].item() > 0.999
    assert layer.hard(scores).item() == 1

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_soft_decision_is_a_distribution_and_hardens" -q

Method implementation

    def hard(self, scores: Tensor) -> Tensor:
        r"""The ``beta -> inf`` hard decision ``argmax(scores)`` (indices)."""
        return torch.argmax(scores, dim=self.axis)

View source ↗

certifiedmethod

certified(self, scores: Tensor, *, eps: float | None=None) -> tuple[Tensor, SelectionCertificate]

Relaxed decision + SelectionCertificate for a single 1-D score vector.

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
epskeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returns: tuple[Tensor, SelectionCertificate]

Executed regression · layer certificate matches functional

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_layer_certificate_matches_functional() -> None:
    layer = DecisionLayer(beta=8.0, eps=0.05)
    scores = torch.tensor([1.0, -0.5, 0.3], dtype=torch.float64)
    soft, cert = layer.certified(scores)
    assert torch.allclose(soft, layer.forward(scores))
    ref = certified_decision(scores.numpy(), beta=8.0, eps=0.05)
    assert cert.argmax == ref.argmax
    assert math.isclose(cert.gap_bound, ref.gap_bound, rel_tol=1e-12)
    assert math.isclose(cert.p_max, ref.p_max, rel_tol=1e-9)

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_layer_certificate_matches_functional" -q

Method implementation

    def certified(
        self, scores: Tensor, *, eps: float | None = None
    ) -> tuple[Tensor, SelectionCertificate]:
        r"""Relaxed decision + :class:`SelectionCertificate` for a single 1-D score vector."""
        soft, cert = certified_argmax(scores, self.beta, eps=self.eps if eps is None else eps)
        return soft, cert

View source ↗

certificatemethod

certificate(self, scores: Tensor, *, eps: float | None=None) -> SelectionCertificate

The SelectionCertificate for a single 1-D score vector (no soft output).

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
epskeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returns: SelectionCertificate

Method implementation

    def certificate(self, scores: Tensor, *, eps: float | None = None) -> SelectionCertificate:
        r"""The :class:`SelectionCertificate` for a single 1-D score vector (no soft output)."""
        _, cert = certified_argmax(scores, self.beta, eps=self.eps if eps is None else eps)
        return cert

View source ↗

certificatesmethod

certificates(self, scores: Tensor, *, eps: float | None=None) -> list[SelectionCertificate]

One SelectionCertificate per row of a (n, N) batch of score vectors.

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
epskeyword-only	float | NoneDefault: None	See the operation contract and implementation below.

Returns: list[SelectionCertificate]

Validation & failure cases

• When s.ndim != 2, raises ValueError(f'certificates expects a 2-D (n, N) batch, got shape {tuple(s.shape)}').

Executed regression · batched certificates and regret

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_batched_certificates_and_regret() -> None:
    layer = DecisionLayer(beta=6.0)
    scores = torch.tensor([[0.1, 0.9, 0.2], [1.2, 0.0, -0.4]], dtype=torch.float64)
    certs = layer.certificates(scores)
    assert len(certs) == 2
    assert [c.argmax for c in certs] == [1, 0]

    rewards = np.array([[0.0, 1.0, 0.5], [2.0, 0.1, 0.0]])
    # perfect predictions -> zero regret
    reg = decision_regret(scores.numpy(), rewards)
    assert np.allclose(reg, 0.0)
    # a wrong prediction on row 0 incurs regret 1.0
    bad = np.array([[3.0, 0.9, 0.2], [1.2, 0.0, -0.4]])
    assert np.isclose(decision_regret(bad, rewards).mean(), 0.5)
    assert best_index(rewards, axis=1).tolist() == [1, 0]

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_batched_certificates_and_regret" -q

Method implementation

    def certificates(
        self, scores: Tensor, *, eps: float | None = None
    ) -> list[SelectionCertificate]:
        r"""One :class:`SelectionCertificate` per row of a ``(n, N)`` batch of score vectors."""
        s = torch.as_tensor(scores)
        if s.ndim != 2:
            raise ValueError(f"certificates expects a 2-D (n, N) batch, got shape {tuple(s.shape)}")
        out: list[SelectionCertificate] = []
        for row in s:
            _, cert = certified_argmax(row, self.beta, eps=self.eps if eps is None else eps)
            out.append(cert)
        return out

View source ↗

extra_reprmethod

extra_repr(self) -> str

See the class contract above and the exact method implementation below.

Returns: str

Method implementation

    def extra_repr(self) -> str:
        return f"beta={self.beta}, eps={self.eps}, axis={self.axis}"

View source ↗

betaattribute

beta = float(beta)

Inverse temperature of the relaxed decision (larger -> sharper / closer to argmax).

View source ↗

epsattribute

eps = None if eps is None else float(eps)

Optional default L^inf perturbation radius for the argmax-stability sub-claim.

View source ↗

axisattribute

axis = int(axis)

The axis the decision is taken over (default the last).

View source ↗

Read the implementation

class DecisionLayer(nn.Module):
    r"""An embeddable certified decision (argmax-over-``N``) layer.

    Parameters
    ----------
    beta:
        Inverse temperature of the relaxed decision (larger -> sharper / closer to argmax).
    eps:
        Optional default ``L^inf`` perturbation radius for the argmax-stability sub-claim.
    axis:
        The axis the decision is taken over (default the last).
    """

    def __init__(self, beta: float = 10.0, *, eps: float | None = None, axis: int = -1) -> None:
        super().__init__()
        if beta <= 0.0:
            raise ValueError(f"beta must be > 0, got {beta}")
        self.beta = float(beta)
        self.eps = None if eps is None else float(eps)
        self.axis = int(axis)

    def forward(self, scores: Tensor) -> Tensor:
        r"""The relaxed decision ``softmax(beta * scores)`` (differentiable)."""
        out: Tensor = soft_argmax(scores, self.beta, axis=self.axis)
        return out

    def hard(self, scores: Tensor) -> Tensor:
        r"""The ``beta -> inf`` hard decision ``argmax(scores)`` (indices)."""
        return torch.argmax(scores, dim=self.axis)

    def certified(
        self, scores: Tensor, *, eps: float | None = None
    ) -> tuple[Tensor, SelectionCertificate]:
        r"""Relaxed decision + :class:`SelectionCertificate` for a single 1-D score vector."""
        soft, cert = certified_argmax(scores, self.beta, eps=self.eps if eps is None else eps)
        return soft, cert

    def certificate(self, scores: Tensor, *, eps: float | None = None) -> SelectionCertificate:
        r"""The :class:`SelectionCertificate` for a single 1-D score vector (no soft output)."""
        _, cert = certified_argmax(scores, self.beta, eps=self.eps if eps is None else eps)
        return cert

    def certificates(
        self, scores: Tensor, *, eps: float | None = None
    ) -> list[SelectionCertificate]:
        r"""One :class:`SelectionCertificate` per row of a ``(n, N)`` batch of score vectors."""
        s = torch.as_tensor(scores)
        if s.ndim != 2:
            raise ValueError(f"certificates expects a 2-D (n, N) batch, got shape {tuple(s.shape)}")
        out: list[SelectionCertificate] = []
        for row in s:
            _, cert = certified_argmax(row, self.beta, eps=self.eps if eps is None else eps)
            out.append(cert)
        return out

    def extra_repr(self) -> str:
        return f"beta={self.beta}, eps={self.eps}, axis={self.axis}"

Related APIsomnibias.struct.torch.certified_argmax↗omnibias.struct.torch.soft_argmax↗

FUNCTION

decision_regret

omnibias.struct.decision.torch.decision_regretSource ↗

decision_regret(scores_hat: Tensor, rewards: Tensor) -> Tensor

Per-sample hard regret of argmax(scores_hat) under true rewards ((n, N) -> (n,)).

Parameters

Name	Type & default	Contract
scores_hatpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rewardspositional-or-keyword	TensorRequired	See the operation contract and implementation below.

ReturnsTensor

Executed regression · layer trains and reduces regret

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_layer_trains_and_reduces_regret() -> None:
    # predict-then-optimize: pick the max-reward option; reward_j is linear in features.
    torch.manual_seed(0)
    rng = np.random.default_rng(0)
    n, d, num_opt = 256, 4, 5
    Wstar = rng.standard_normal((d, num_opt))
    X = rng.standard_normal((n, d))
    rewards = X @ Wstar  # (n, num_opt) true option rewards
    Xt = torch.tensor(X, dtype=torch.float64)
    Rt = torch.tensor(rewards, dtype=torch.float64)

    model = torch.nn.Linear(d, num_opt, dtype=torch.float64)
    layer = DecisionLayer(beta=8.0)
    opt = torch.optim.Adam(model.parameters(), lr=5e-2)

    def mean_regret() -> float:
        with torch.no_grad():
            return float(decision_regret_torch(model(Xt), Rt).mean().item())

    init_regret = mean_regret()
    for _ in range(150):
        opt.zero_grad()
        p = layer(model(Xt))  # relaxed decision (differentiable)
        loss = -expected_reward(p, Rt).mean()  # maximize expected realized reward
        loss.backward()
        opt.step()
    final_regret = mean_regret()

    # training through the differentiable decision drives the hard regret down
    assert final_regret < 0.5 * init_regret

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_layer_trains_and_reduces_regret" -q

Read the implementation

def decision_regret(scores_hat: Tensor, rewards: Tensor) -> Tensor:
    r"""Per-sample hard regret of ``argmax(scores_hat)`` under true ``rewards`` (``(n, N)`` -> ``(n,)``)."""
    picked = torch.argmax(scores_hat, dim=-1)
    realized = rewards.gather(-1, picked.unsqueeze(-1)).squeeze(-1)
    return rewards.max(dim=-1).values - realized

FUNCTION

expected_reward

omnibias.struct.decision.torch.expected_rewardSource ↗

expected_reward(p: Tensor, rewards: Tensor, *, axis: int=-1) -> Tensor

The differentiable decision objective E_p[reward] = <p, rewards> along axis.

Parameters

Name	Type & default	Contract
ppositional-or-keyword	TensorRequired	See the operation contract and implementation below.
rewardspositional-or-keyword	TensorRequired	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

ReturnsTensor

Executed regression · layer trains and reduces regret

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_layer_trains_and_reduces_regret() -> None:
    # predict-then-optimize: pick the max-reward option; reward_j is linear in features.
    torch.manual_seed(0)
    rng = np.random.default_rng(0)
    n, d, num_opt = 256, 4, 5
    Wstar = rng.standard_normal((d, num_opt))
    X = rng.standard_normal((n, d))
    rewards = X @ Wstar  # (n, num_opt) true option rewards
    Xt = torch.tensor(X, dtype=torch.float64)
    Rt = torch.tensor(rewards, dtype=torch.float64)

    model = torch.nn.Linear(d, num_opt, dtype=torch.float64)
    layer = DecisionLayer(beta=8.0)
    opt = torch.optim.Adam(model.parameters(), lr=5e-2)

    def mean_regret() -> float:
        with torch.no_grad():
            return float(decision_regret_torch(model(Xt), Rt).mean().item())

    init_regret = mean_regret()
    for _ in range(150):
        opt.zero_grad()
        p = layer(model(Xt))  # relaxed decision (differentiable)
        loss = -expected_reward(p, Rt).mean()  # maximize expected realized reward
        loss.backward()
        opt.step()
    final_regret = mean_regret()

    # training through the differentiable decision drives the hard regret down
    assert final_regret < 0.5 * init_regret

python -m pytest "packages/omnibias-struct/tests/test_decision_layer.py::test_layer_trains_and_reduces_regret" -q

Read the implementation

def expected_reward(p: Tensor, rewards: Tensor, *, axis: int = -1) -> Tensor:
    r"""The differentiable decision objective ``E_p[reward] = <p, rewards>`` along ``axis``."""
    return (p * rewards).sum(dim=axis)

FUNCTION

soft_decision

omnibias.struct.decision.torch.soft_decisionSource ↗

soft_decision(scores: Tensor, beta: float, *, axis: int=-1) -> Tensor

The relaxed decision softmax(beta * scores) (-> one-hot argmax as beta -> inf).

Parameters

Name	Type & default	Contract
scorespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
betapositional-or-keyword	floatRequired	See the operation contract and implementation below.
axiskeyword-only	intDefault: -1	See the operation contract and implementation below.

ReturnsTensor

Read the implementation

def soft_decision(scores: Tensor, beta: float, *, axis: int = -1) -> Tensor:
    r"""The relaxed decision ``softmax(beta * scores)`` (``-> one-hot argmax`` as ``beta -> inf``)."""
    out: Tensor = soft_argmax(scores, beta, axis=axis)
    return out

Related APIsomnibias.struct.torch.soft_argmax↗

← Package overviewRunnable learning examples →