API REFERENCE/omnibias-core

omnibias.core.realization.transition

0.5.0rc1Release candidatePython >=3.10Apache-2.0

Coefficient recurrences, activation contracts, jet algebra, and verified arithmetic shared by the numerical backends.

import omnibias.core.realization.transition

Module contract & mathematical conventions

Portable, versioned structural proposals. Acceptance is a downstream concern.

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

SlotUpdate

omnibias.core.realization.transition.SlotUpdateSource ↗

class SlotUpdate

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
namedataclass field	strRequired	See the operation contract and implementation below.
indexdataclass field	intRequired	See the operation contract and implementation below.
valuedataclass field	float | tuple[float, ...]Required	See the operation contract and implementation below.

Fields & methods

namefield

name: str

Declared str field on SlotUpdate. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

indexfield

index: int

Declared int field on SlotUpdate. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

valuefield

value: float | tuple[float, ...]

Declared float | tuple[float, ...] field on SlotUpdate. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

Read the implementation

class SlotUpdate:
    name: str
    index: int
    value: float | tuple[float, ...]

CLASS

TransitionProposal

omnibias.core.realization.transition.TransitionProposalSource ↗

class TransitionProposal

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
source_digestdataclass field	strRequired	See the operation contract and implementation below.
source_versiondataclass field	intRequired	See the operation contract and implementation below.
slot_idsdataclass field	tuple[int, ...]Required	See the operation contract and implementation below.
updatesdataclass field	tuple[SlotUpdate, ...]Required	See the operation contract and implementation below.
kinddataclass field	strRequired	See the operation contract and implementation below.
coordinate_identitydataclass field	strRequired	See the operation contract and implementation below.
error_budgetdataclass field	floatDefault: 0.0	See the operation contract and implementation below.
schema_versiondataclass field	intDefault: 1	See the operation contract and implementation below.

Regression example · cluster row transition and error gated derivative

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_cluster_row_transition_and_error_gated_derivative() -> None:
    packs = [RefinedPack(0, v, w) for v, w in [(-0.01, 0.7), (0.0, -0.2), (0.02, 0.5)]]
    bank = jb.init_confluent_bank(packs, max_packs=4, max_order=4)
    torch_bank = tb.ConfluentPackBank(packs, max_packs=4, max_order=4, dtype=torch.float64)
    proposal = bank.propose_cluster((0, 1, 2), order=3, error_budget=1e-7)
    assert proposal == torch_bank.propose_cluster((0, 1, 2), order=3, error_budget=1e-7)
    assert TransitionProposal.from_json(proposal.to_json()) == proposal
    candidate, _, accepted = jb.commit_transition(bank, proposal, moments(bank), between_steps=True,
                                                  acceptance=lambda _b, _p: True)
    assert accepted and int(candidate.modes[0]) == 2
    assert candidate.moment_coefficients.shape == (4, 5)
    assert not bool(candidate.active[1]) and not bool(candidate.active[2])
    np.testing.assert_allclose(np.asarray(candidate.forward(jnp.array([-0.2, 0.4]))),
                               np.asarray(bank.forward(jnp.array([-0.2, 0.4]))), atol=2e-9)
    _, pair = pair_banks()
    pair, state, _ = jb.commit_transition(pair, pair.propose_pair(0, 1), moments(pair), between_steps=True,
                                          acceptance=lambda _b, _p: True)
    derivative = pair.propose_derivative(0, error_budget=0)
    same, same_state, accepted = jb.commit_transition(pair, derivative, state, between_steps=True,
                                                     acceptance=lambda _b, p: p.error_budget > 0)
    assert not accepted and same is pair and same_state is state
    assert float(pair.rho[0]) > 0

python -m pytest tests/test_confluent_bank_parity.py -k test_cluster_row_transition_and_error_gated_derivative -q

Regression example · portable proposals and functional transaction match torch

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_portable_proposals_and_functional_transaction_match_torch() -> None:
    torch_bank, bank = pair_banks()
    assert bank.portable_snapshot() == torch_bank.portable_snapshot()
    proposal = bank.propose_pair(0, 1)
    assert proposal == torch_bank.propose_pair(0, 1)
    assert TransitionProposal.from_json(proposal.to_json()) == proposal
    state = moments(bank)
    candidate, next_state, accepted = jb.commit_transition(bank, proposal, state, between_steps=True,
                                                          acceptance=lambda _b, _p: True)
    assert accepted and int(candidate.transition_version) == 1
    assert int(bank.transition_version) == 0 and bool(bank.active[1])
    assert not bool(candidate.active[1])
    for kind, entries in next_state.moments.items():
        for name, value in entries.items():
            np.testing.assert_array_equal(np.asarray(value[:2]), np.zeros(value[:2].shape))
            np.testing.assert_array_equal(np.asarray(value[2:]), np.asarray(state.moments[kind][name][2:]))
    assert int(next_state.step) == 7
    optimizer = torch.optim.Adam(torch_bank.parameters())
    for parameter in torch_bank.parameters():
        optimizer.state[parameter] = {"step": torch.tensor(7.0), "exp_avg": torch.ones_like(parameter),
                                      "exp_avg_sq": 2 * torch.ones_like(parameter)}
    assert tb.commit_transition(torch_bank, proposal, optimizer, between_steps=True,
                                acceptance=lambda _b, _p: True)
    assert candidate.portable_snapshot() == torch_bank.portable_snapshot()
    for name, parameter in torch_bank.named_parameters():
        np.testing.assert_array_equal(optimizer.state[parameter]["exp_avg"].numpy(),
                                      np.asarray(next_state.moments["first"][name]))
    x = jnp.linspace(-0.5, 0.5, 9)
    np.testing.assert_allclose(np.asarray(candidate.forward(x)), np.asarray(bank.forward(x)), atol=3e-16)
    np.testing.assert_allclose(torch_bank(torch.tensor(np.asarray(x))).detach().numpy(),
                               np.asarray(candidate.forward(x)), atol=3e-16)
    with pytest.raises(ValueError, match="stale"):
        jb.commit_transition(candidate, proposal, next_state, between_steps=True,
                              acceptance=lambda _b, _p: True)

python -m pytest tests/test_confluent_bank_parity.py -k test_portable_proposals_and_functional_transaction_match_torch -q

Fields & methods

source_digestfield

source_digest: str

Declared str field on TransitionProposal. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

source_versionfield

source_version: int

Declared int field on TransitionProposal. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

slot_idsfield

slot_ids: tuple[int, ...]

Declared tuple[int, ...] field on TransitionProposal. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

updatesfield

updates: tuple[SlotUpdate, ...]

Declared tuple[SlotUpdate, ...] field on TransitionProposal. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

kindfield

kind: str

Declared str field on TransitionProposal. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

coordinate_identityfield

coordinate_identity: str

Declared str field on TransitionProposal. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

error_budgetfield

error_budget: float = 0.0

Declared float field on TransitionProposal. Default expression: 0.0.

View source ↗

schema_versionfield

schema_version: int = 1

Declared int field on TransitionProposal. Default expression: 1.

View source ↗

__post_init__method

__post_init__(self) -> None

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

Returns: None

Validation & failure cases

• When self.error_budget < 0 or self.schema_version != 1, raises ValueError('invalid transition budget or schema').
• When len(set(self.slot_ids)) != len(self.slot_ids), raises ValueError('slot IDs must be unique').

Method implementation

    def __post_init__(self) -> None:
        if self.error_budget < 0 or self.schema_version != 1:
            raise ValueError('invalid transition budget or schema')
        if len(set(self.slot_ids)) != len(self.slot_ids):
            raise ValueError('slot IDs must be unique')

View source ↗

to_jsonmethod

to_json(self) -> str

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

Returns: str

Method implementation

    def to_json(self) -> str:
        return json.dumps(asdict(self), sort_keys=True, separators=(',', ':'), allow_nan=False)

View source ↗

from_jsonmethod

from_json(cls, text: str) -> TransitionProposal

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

Parameters

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

Returns: TransitionProposal

Regression example · cluster row transition and error gated derivative

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_cluster_row_transition_and_error_gated_derivative() -> None:
    packs = [RefinedPack(0, v, w) for v, w in [(-0.01, 0.7), (0.0, -0.2), (0.02, 0.5)]]
    bank = jb.init_confluent_bank(packs, max_packs=4, max_order=4)
    torch_bank = tb.ConfluentPackBank(packs, max_packs=4, max_order=4, dtype=torch.float64)
    proposal = bank.propose_cluster((0, 1, 2), order=3, error_budget=1e-7)
    assert proposal == torch_bank.propose_cluster((0, 1, 2), order=3, error_budget=1e-7)
    assert TransitionProposal.from_json(proposal.to_json()) == proposal
    candidate, _, accepted = jb.commit_transition(bank, proposal, moments(bank), between_steps=True,
                                                  acceptance=lambda _b, _p: True)
    assert accepted and int(candidate.modes[0]) == 2
    assert candidate.moment_coefficients.shape == (4, 5)
    assert not bool(candidate.active[1]) and not bool(candidate.active[2])
    np.testing.assert_allclose(np.asarray(candidate.forward(jnp.array([-0.2, 0.4]))),
                               np.asarray(bank.forward(jnp.array([-0.2, 0.4]))), atol=2e-9)
    _, pair = pair_banks()
    pair, state, _ = jb.commit_transition(pair, pair.propose_pair(0, 1), moments(pair), between_steps=True,
                                          acceptance=lambda _b, _p: True)
    derivative = pair.propose_derivative(0, error_budget=0)
    same, same_state, accepted = jb.commit_transition(pair, derivative, state, between_steps=True,
                                                     acceptance=lambda _b, p: p.error_budget > 0)
    assert not accepted and same is pair and same_state is state
    assert float(pair.rho[0]) > 0

python -m pytest tests/test_confluent_bank_parity.py -k test_cluster_row_transition_and_error_gated_derivative -q

Method implementation

    def from_json(cls, text: str) -> TransitionProposal:
        data: dict[str, Any] = json.loads(text)
        data['slot_ids'] = tuple(data['slot_ids'])
        data['updates'] = tuple(SlotUpdate(item['name'], item['index'],
                                         tuple(item['value']) if isinstance(item['value'], list)
                                         else item['value']) for item in data['updates'])
        return cls(**data)

View source ↗

Read the implementation

class TransitionProposal:
    source_digest: str
    source_version: int
    slot_ids: tuple[int, ...]
    updates: tuple[SlotUpdate, ...]
    kind: str
    coordinate_identity: str
    error_budget: float = 0.0
    schema_version: int = 1

    def __post_init__(self) -> None:
        if self.error_budget < 0 or self.schema_version != 1:
            raise ValueError('invalid transition budget or schema')
        if len(set(self.slot_ids)) != len(self.slot_ids):
            raise ValueError('slot IDs must be unique')

    def to_json(self) -> str:
        return json.dumps(asdict(self), sort_keys=True, separators=(',', ':'), allow_nan=False)

    @classmethod
    def from_json(cls, text: str) -> TransitionProposal:
        data: dict[str, Any] = json.loads(text)
        data['slot_ids'] = tuple(data['slot_ids'])
        data['updates'] = tuple(SlotUpdate(item['name'], item['index'],
                                         tuple(item['value']) if isinstance(item['value'], list)
                                         else item['value']) for item in data['updates'])
        return cls(**data)

Related APIsomnibias.core.realization.transition.SlotUpdate↗

FUNCTION

snapshot_digest

omnibias.core.realization.transition.snapshot_digestSource ↗

snapshot_digest(snapshot: dict[str, Any]) -> str

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
snapshotpositional-or-keyword	dict[str, Any]Required	See the operation contract and implementation below.

Returnsstr

Read the implementation

def snapshot_digest(snapshot: dict[str, Any]) -> str:
    return sha256(json.dumps(snapshot, sort_keys=True, separators=(',', ':'),
                             allow_nan=False).encode()).hexdigest()

← Package overviewRunnable learning examples →