API REFERENCE/omnibias-torch

omnibias.torch.refine

0.5.0rc1Release candidatePython >=3.10Apache-2.0

PyTorch activation derivatives, trainable operator layers, network jets, and direct differential-operator contractions.

import omnibias.torch.refine

Module contract & mathematical conventions

Adaptive pack refinement (theory 03-13), PyTorch twin.

AdaptivePackBank is a pre-allocated bank of tempered packs

c * alpha^n sigma^(n)(alpha (z - center))

Birth and *p*-type growth write a new slot with c = 0 (bit-identical output). Death deactivates a slot and reports ~omnibias.core.refine.RefineReport.death_perturbation.

**Optimizer-state policy** (same as GrowableOperatorMultiBiasUnit): slots are allocated up to max_packs. Inactive slots do not enter the forward sum, so they receive no gradient and Adam-style moments stay at zero until first activation. Death leaves whatever moment state the slot had; a later reuse of that slot inherits it. Callers who hold an optimizer may zero that slot's moments on death. Do not reconstruct the module mid-step.

Literal-K growth of a single OMBU still goes through GrowableOperatorMultiBiasUnit.grow via grow_ombu. That path is torch-only (there is no jax GrowableOMBU). Pack-bank *p*-type growth is the collapsed sibling and is the jax-parity path.

The refine driver is eager. Do not wrap it in torch.compile.

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
GrowStrategy	omnibias.torch.GrowStrategy
Indicator	omnibias.core.Indicator
RefinePolicy	omnibias.core.RefinePolicy
RefineReport	omnibias.core.RefineReport
RefinedPack	omnibias.core.RefinedPack

CLASS

AdaptivePackBank

omnibias.torch.AdaptivePackBankSource ↗

class AdaptivePackBank(nn.Module)

Pre-allocated tempered pack bank with birth / growth / death.

Parameters

packs: Initial live packs. Length must be <= max_packs. max_packs: Slot budget. Optimizer moments for unused slots stay at zero. max_order: Highest derivative order a slot may hold (static loop bound). base: Activation with a closed-form fastpath.

Parameters

Name	Type & default	Contract
packspositional-or-keyword	Sequence[RefinedPack]Required	Initial live packs. Length must be <= max_packs.
max_packskeyword-only	intDefault: 16	Slot budget. Optimizer moments for unused slots stay at zero.
max_orderkeyword-only	intDefault: 8	Highest derivative order a slot may hold (static loop bound).
basekeyword-only	str | ActivationSpec[Tensor]Default: 'exp'	Activation with a closed-form fastpath.
learnable_centerskeyword-only	boolDefault: True	See the operation contract and implementation below.
learnable_weightskeyword-only	boolDefault: True	See the operation contract and implementation below.
learnable_scaleskeyword-only	boolDefault: True	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Regression example · budget stability with death

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_budget_stability_with_death() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 17)
    policy = RefinePolicy(
        indicator=Indicator.RESIDUAL,
        birth_threshold=0.5,
        death_threshold=0.2,
        min_age=2,
        hysteresis=1.0,
        max_packs=6,
        min_scale_ratio=1.0,
    )
    counts: list[int] = []
    for step in range(24):
        amp = 0.8 if step % 4 == 0 else 0.0

        def residual(x: torch.Tensor, _amp: float = amp) -> torch.Tensor:
            return torch.full_like(x, _amp)

        refine(bank, residual, probes, policy, step=step)
        counts.append(bank.n_active)
    assert max(counts) <= 6
    assert counts[-1] <= counts[0] + 2

python -m pytest packages/omnibias-torch/tests/test_refine.py -k test_budget_stability_with_death -q

Fields & methods

act_specfield

act_spec: ActivationSpec[Tensor]

Declared ActivationSpec[Tensor] field on AdaptivePackBank. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

centersfield

centers: Tensor

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

View source ↗

weightsfield

weights: Tensor

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

View source ↗

scalesfield

scales: Tensor

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

View source ↗

ordersfield

orders: Tensor

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

View source ↗

agesfield

ages: Tensor

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

View source ↗

activefield

active: Tensor

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

View source ↗

__init__method

__init__(self, packs: Sequence[RefinedPack], *, max_packs: int=16, max_order: int=8, base: str | ActivationSpec[Tensor]='exp', learnable_centers: bool=True, learnable_weights: bool=True, learnable_scales: bool=True, dtype: torch.dtype | None=None) -> None

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

Parameters

Name	Type & default	Contract
packspositional-or-keyword	Sequence[RefinedPack]Required	See the operation contract and implementation below.
max_packskeyword-only	intDefault: 16	See the operation contract and implementation below.
max_orderkeyword-only	intDefault: 8	See the operation contract and implementation below.
basekeyword-only	str | ActivationSpec[Tensor]Default: 'exp'	See the operation contract and implementation below.
learnable_centerskeyword-only	boolDefault: True	See the operation contract and implementation below.
learnable_weightskeyword-only	boolDefault: True	See the operation contract and implementation below.
learnable_scaleskeyword-only	boolDefault: True	See the operation contract and implementation below.
dtypekeyword-only	torch.dtype | NoneDefault: None	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When int(max_packs) < 1, raises ValueError(f'max_packs must be >= 1, got {max_packs}').
• When int(max_order) < 0, raises ValueError(f'max_order must be >= 0, got {max_order}').
• When len(live) > int(max_packs), raises ValueError(f'{len(live)} initial packs exceed max_packs={max_packs}').
• When pack.order > int(max_order), raises ValueError(f'pack order {pack.order} exceeds max_order={max_order}').
• When self.act_spec.fastpath is None, raises NotImplementedError(f'Activation {self.act_spec.name!r} has no closed-form derivative kernel').

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Method implementation

    def __init__(
        self,
        packs: Sequence[RefinedPack],
        *,
        max_packs: int = 16,
        max_order: int = 8,
        base: str | ActivationSpec[Tensor] = "exp",
        learnable_centers: bool = True,
        learnable_weights: bool = True,
        learnable_scales: bool = True,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if int(max_packs) < 1:
            raise ValueError(f"max_packs must be >= 1, got {max_packs}")
        if int(max_order) < 0:
            raise ValueError(f"max_order must be >= 0, got {max_order}")
        live = tuple(packs)
        if len(live) > int(max_packs):
            raise ValueError(f"{len(live)} initial packs exceed max_packs={max_packs}")
        for pack in live:
            if pack.order > int(max_order):
                raise ValueError(
                    f"pack order {pack.order} exceeds max_order={max_order}"
                )
        self.max_packs = int(max_packs)
        self.max_order = int(max_order)
        self.act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        if self.act_spec.fastpath is None:
            raise NotImplementedError(
                f"Activation {self.act_spec.name!r} has no closed-form derivative kernel"
            )
        dt = torch.get_default_dtype() if dtype is None else dtype
        centers0 = torch.zeros(self.max_packs, dtype=dt)
        weights0 = torch.zeros(self.max_packs, dtype=dt)
        scales0 = torch.ones(self.max_packs, dtype=dt)
        orders0 = torch.zeros(self.max_packs, dtype=torch.long)
        ages0 = torch.zeros(self.max_packs, dtype=torch.long)
        active0 = torch.zeros(self.max_packs, dtype=torch.bool)
        for i, pack in enumerate(live):
            centers0[i] = pack.center
            weights0[i] = pack.weight
            scales0[i] = pack.scale
            orders0[i] = pack.order
            ages0[i] = pack.age
            active0[i] = True
        if learnable_centers:
            self.centers = nn.Parameter(centers0)
        else:
            self.register_buffer("centers", centers0)
        if learnable_weights:
            self.weights = nn.Parameter(weights0)
        else:
            self.register_buffer("weights", weights0)
        if learnable_scales:
            self.scales = nn.Parameter(scales0)
        else:
            self.register_buffer("scales", scales0)
        self.register_buffer("orders", orders0)
        self.register_buffer("ages", ages0)
        self.register_buffer("active", active0)

View source ↗

n_activeproperty

n_active(self) -> int

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

Returns: int

Executed regression · budget stability with death

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_budget_stability_with_death() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 17)
    policy = RefinePolicy(
        indicator=Indicator.RESIDUAL,
        birth_threshold=0.5,
        death_threshold=0.2,
        min_age=2,
        hysteresis=1.0,
        max_packs=6,
        min_scale_ratio=1.0,
    )
    counts: list[int] = []
    for step in range(24):
        amp = 0.8 if step % 4 == 0 else 0.0

        def residual(x: torch.Tensor, _amp: float = amp) -> torch.Tensor:
            return torch.full_like(x, _amp)

        refine(bank, residual, probes, policy, step=step)
        counts.append(bank.n_active)
    assert max(counts) <= 6
    assert counts[-1] <= counts[0] + 2

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_budget_stability_with_death" -q

Method implementation

    def n_active(self) -> int:
        return int(self.active.sum().item())

View source ↗

active_packsmethod

active_packs(self) -> list[RefinedPack]

Live packs in slot order (stable, used as the refine algebra list).

Returns: list[RefinedPack]

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Method implementation

    def active_packs(self) -> list[RefinedPack]:
        """Live packs in slot order (stable, used as the refine algebra list)."""
        out: list[RefinedPack] = []
        for i in range(self.max_packs):
            if bool(self.active[i].item()):
                out.append(self._pack_at(i))
        return out

View source ↗

forwardmethod

forward(self, z: Tensor) -> Tensor

1-D field z → same leading shape.

Parameters

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

Returns: Tensor

Validation & failure cases

• When self.act_spec.fastpath is None, raises NotImplementedError(f'Activation {self.act_spec.name!r} has no closed-form derivative kernel').

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Method implementation

    def forward(self, z: Tensor) -> Tensor:
        """1-D field ``z`` → same leading shape."""
        flat = _as_1d(z)
        if self.act_spec.fastpath is None:
            raise NotImplementedError(
                f"Activation {self.act_spec.name!r} has no closed-form derivative kernel"
            )
        out = torch.zeros_like(flat)
        for i in range(self.max_packs):
            if not bool(self.active[i].item()):
                continue
            n = int(self.orders[i].item())
            out = out + pack_response(
                flat,
                self.centers[i],
                self.weights[i],
                self.scales[i],
                n,
                self.act_spec,
            )
        return out.reshape(z.shape) if z.ndim != 1 and z.shape[-1] == 1 else out

View source ↗

pack_term_normsmethod

pack_term_norms(self, z: Tensor) -> list[float]

L2 norms of each live pack's term on z, in slot order.

Parameters

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

Returns: list[float]

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Method implementation

    def pack_term_norms(self, z: Tensor) -> list[float]:
        """L2 norms of each live pack's term on ``z``, in slot order."""
        flat = _as_1d(z)
        norms: list[float] = []
        for i in range(self.max_packs):
            if not bool(self.active[i].item()):
                continue
            n = int(self.orders[i].item())
            term = pack_response(
                flat,
                self.centers[i].detach(),
                self.weights[i].detach(),
                self.scales[i].detach(),
                n,
                self.act_spec,
            )
            norms.append(float(term.square().sum().sqrt().item()))
        return norms

View source ↗

sync_frommethod

sync_from(self, packs: Sequence[RefinedPack]) -> None

Write a refined list back onto slots, keeping live packs in place.

Parameters

Name	Type & default	Contract
packspositional-or-keyword	Sequence[RefinedPack]Required	See the operation contract and implementation below.

Returns: None

Validation & failure cases

• When free is None, raises RuntimeError(f'max_packs={self.max_packs} exhausted').

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Method implementation

    def sync_from(self, packs: Sequence[RefinedPack]) -> None:
        """Write a refined list back onto slots, keeping live packs in place."""
        occupied: dict[tuple[int, float, float, float], int] = {}
        for i in range(self.max_packs):
            if bool(self.active[i].item()):
                occupied[self._key(self._pack_at(i))] = i
        assigned: set[int] = set()
        for pack in packs:
            key = self._key(pack)
            slot = occupied.get(key)
            if slot is not None and slot not in assigned:
                with torch.no_grad():
                    self.ages[slot] = pack.age
                assigned.add(slot)
                continue
            free = next((i for i in range(self.max_packs) if i not in assigned), None)
            if free is None:
                raise RuntimeError(f"max_packs={self.max_packs} exhausted")
            self._write_slot(free, pack)
            assigned.add(free)
        for i in range(self.max_packs):
            if i not in assigned:
                self._deactivate(i)

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"n_active={self.n_active}/{self.max_packs}, "
            f"base={self.act_spec.name!r}, max_order={self.max_order}"
        )

View source ↗

max_packsattribute

max_packs = int(max_packs)

Slot budget. Optimizer moments for unused slots stay at zero.

View source ↗

max_orderattribute

max_order = int(max_order)

Highest derivative order a slot may hold (static loop bound).

View source ↗

Read the implementation

class AdaptivePackBank(nn.Module):
    """Pre-allocated tempered pack bank with birth / growth / death.

    Parameters
    ----------
    packs:
        Initial live packs. Length must be ``<= max_packs``.
    max_packs:
        Slot budget. Optimizer moments for unused slots stay at zero.
    max_order:
        Highest derivative order a slot may hold (static loop bound).
    base:
        Activation with a closed-form fastpath.
    """

    act_spec: ActivationSpec[Tensor]
    centers: Tensor
    weights: Tensor
    scales: Tensor
    orders: Tensor
    ages: Tensor
    active: Tensor

    def __init__(
        self,
        packs: Sequence[RefinedPack],
        *,
        max_packs: int = 16,
        max_order: int = 8,
        base: str | ActivationSpec[Tensor] = "exp",
        learnable_centers: bool = True,
        learnable_weights: bool = True,
        learnable_scales: bool = True,
        dtype: torch.dtype | None = None,
    ) -> None:
        super().__init__()
        if int(max_packs) < 1:
            raise ValueError(f"max_packs must be >= 1, got {max_packs}")
        if int(max_order) < 0:
            raise ValueError(f"max_order must be >= 0, got {max_order}")
        live = tuple(packs)
        if len(live) > int(max_packs):
            raise ValueError(f"{len(live)} initial packs exceed max_packs={max_packs}")
        for pack in live:
            if pack.order > int(max_order):
                raise ValueError(
                    f"pack order {pack.order} exceeds max_order={max_order}"
                )
        self.max_packs = int(max_packs)
        self.max_order = int(max_order)
        self.act_spec = base if isinstance(base, ActivationSpec) else get_activation(base)
        if self.act_spec.fastpath is None:
            raise NotImplementedError(
                f"Activation {self.act_spec.name!r} has no closed-form derivative kernel"
            )
        dt = torch.get_default_dtype() if dtype is None else dtype
        centers0 = torch.zeros(self.max_packs, dtype=dt)
        weights0 = torch.zeros(self.max_packs, dtype=dt)
        scales0 = torch.ones(self.max_packs, dtype=dt)
        orders0 = torch.zeros(self.max_packs, dtype=torch.long)
        ages0 = torch.zeros(self.max_packs, dtype=torch.long)
        active0 = torch.zeros(self.max_packs, dtype=torch.bool)
        for i, pack in enumerate(live):
            centers0[i] = pack.center
            weights0[i] = pack.weight
            scales0[i] = pack.scale
            orders0[i] = pack.order
            ages0[i] = pack.age
            active0[i] = True
        if learnable_centers:
            self.centers = nn.Parameter(centers0)
        else:
            self.register_buffer("centers", centers0)
        if learnable_weights:
            self.weights = nn.Parameter(weights0)
        else:
            self.register_buffer("weights", weights0)
        if learnable_scales:
            self.scales = nn.Parameter(scales0)
        else:
            self.register_buffer("scales", scales0)
        self.register_buffer("orders", orders0)
        self.register_buffer("ages", ages0)
        self.register_buffer("active", active0)

    @property
    def n_active(self) -> int:
        return int(self.active.sum().item())

    def active_packs(self) -> list[RefinedPack]:
        """Live packs in slot order (stable, used as the refine algebra list)."""
        out: list[RefinedPack] = []
        for i in range(self.max_packs):
            if bool(self.active[i].item()):
                out.append(self._pack_at(i))
        return out

    def _pack_at(self, slot: int) -> RefinedPack:
        return RefinedPack(
            order=int(self.orders[slot].item()),
            center=float(self.centers[slot].detach()),
            weight=float(self.weights[slot].detach()),
            scale=float(self.scales[slot].detach()),
            age=int(self.ages[slot].item()),
        )

    def _key(self, pack: RefinedPack) -> tuple[int, float, float, float]:
        return (pack.order, pack.center, pack.scale, pack.weight)

    def forward(self, z: Tensor) -> Tensor:
        """1-D field ``z`` → same leading shape."""
        flat = _as_1d(z)
        if self.act_spec.fastpath is None:
            raise NotImplementedError(
                f"Activation {self.act_spec.name!r} has no closed-form derivative kernel"
            )
        out = torch.zeros_like(flat)
        for i in range(self.max_packs):
            if not bool(self.active[i].item()):
                continue
            n = int(self.orders[i].item())
            out = out + pack_response(
                flat,
                self.centers[i],
                self.weights[i],
                self.scales[i],
                n,
                self.act_spec,
            )
        return out.reshape(z.shape) if z.ndim != 1 and z.shape[-1] == 1 else out

    def pack_term_norms(self, z: Tensor) -> list[float]:
        """L2 norms of each live pack's term on ``z``, in slot order."""
        flat = _as_1d(z)
        norms: list[float] = []
        for i in range(self.max_packs):
            if not bool(self.active[i].item()):
                continue
            n = int(self.orders[i].item())
            term = pack_response(
                flat,
                self.centers[i].detach(),
                self.weights[i].detach(),
                self.scales[i].detach(),
                n,
                self.act_spec,
            )
            norms.append(float(term.square().sum().sqrt().item()))
        return norms

    def _write_slot(self, slot: int, pack: RefinedPack) -> None:
        with torch.no_grad():
            self.centers[slot] = pack.center
            self.weights[slot] = pack.weight
            self.scales[slot] = pack.scale
            self.orders[slot] = pack.order
            self.ages[slot] = pack.age
            self.active[slot] = True

    def _deactivate(self, slot: int) -> None:
        with torch.no_grad():
            self.active[slot] = False
            self.weights[slot] = 0.0
            self.ages[slot] = 0

    def sync_from(self, packs: Sequence[RefinedPack]) -> None:
        """Write a refined list back onto slots, keeping live packs in place."""
        occupied: dict[tuple[int, float, float, float], int] = {}
        for i in range(self.max_packs):
            if bool(self.active[i].item()):
                occupied[self._key(self._pack_at(i))] = i
        assigned: set[int] = set()
        for pack in packs:
            key = self._key(pack)
            slot = occupied.get(key)
            if slot is not None and slot not in assigned:
                with torch.no_grad():
                    self.ages[slot] = pack.age
                assigned.add(slot)
                continue
            free = next((i for i in range(self.max_packs) if i not in assigned), None)
            if free is None:
                raise RuntimeError(f"max_packs={self.max_packs} exhausted")
            self._write_slot(free, pack)
            assigned.add(free)
        for i in range(self.max_packs):
            if i not in assigned:
                self._deactivate(i)

    def extra_repr(self) -> str:
        return (
            f"n_active={self.n_active}/{self.max_packs}, "
            f"base={self.act_spec.name!r}, max_order={self.max_order}"
        )

Related APIsomnibias.core.RefinedPack↗omnibias.torch.get_activation↗omnibias.torch.refine.pack_response↗

Other public import paths

• omnibias.torch.AdaptivePackBank
• omnibias.torch.refine.AdaptivePackBank

FUNCTION

grow_ombu

omnibias.torch.grow_ombuSource ↗

grow_ombu(unit: GrowableOperatorMultiBiasUnit, strategy: GrowStrategy='pair', **kwargs: object) -> int

Delegate *p*-type growth to the existing growable OMBU. Do not reimplement it.

Parameters

Name	Type & default	Contract
unitpositional-or-keyword	GrowableOperatorMultiBiasUnitRequired	See the operation contract and implementation below.
strategypositional-or-keyword	GrowStrategyDefault: 'pair'	See the operation contract and implementation below.
kwargsvariadic keyword	objectRequired	See the operation contract and implementation below.

Returnsint

Executed regression · grow ombu reuses existing zero perturbation

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_grow_ombu_reuses_existing_zero_perturbation() -> None:
    unit = GrowableOperatorMultiBiasUnit(num_channels=1, init_K=1, K_max=8)
    z = torch.linspace(-1.0, 1.0, 17).unsqueeze(-1)
    before = unit(z).clone()
    added = grow_ombu(unit, "pair")
    assert added == 2
    assert torch.equal(before, unit(z))

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_grow_ombu_reuses_existing_zero_perturbation" -q

Read the implementation

def grow_ombu(
    unit: GrowableOperatorMultiBiasUnit,
    strategy: GrowStrategy = "pair",
    **kwargs: object,
) -> int:
    """Delegate *p*-type growth to the existing growable OMBU. Do not reimplement it."""
    return unit.grow(strategy, **kwargs)  # type: ignore[arg-type]

Other public import paths

• omnibias.torch.grow_ombu
• omnibias.torch.refine.grow_ombu

FUNCTION

pack_response

omnibias.torch.refine.pack_responseSource ↗

pack_response(z: Tensor, center: Tensor, weight: Tensor, scale: Tensor, order: int, spec: ActivationSpec[Tensor]) -> Tensor

c * alpha^n sigma^(n)(alpha (z - center)).

Parameters

Name	Type & default	Contract
zpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
centerpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
weightpositional-or-keyword	TensorRequired	See the operation contract and implementation below.
scalepositional-or-keyword	TensorRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.
specpositional-or-keyword	ActivationSpec[Tensor]Required	See the operation contract and implementation below.

ReturnsTensor

Explicit validation & failure cases

• When spec.fastpath is None, raises NotImplementedError(f'Activation {spec.name!r} has no closed-form derivative kernel').
• When order < 0, raises ValueError(f'order must be >= 0, got {order}').

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

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Read the implementation

def pack_response(
    z: Tensor,
    center: Tensor,
    weight: Tensor,
    scale: Tensor,
    order: int,
    spec: ActivationSpec[Tensor],
) -> Tensor:
    """``c * alpha^n sigma^(n)(alpha (z - center))``."""
    if spec.fastpath is None:
        raise NotImplementedError(
            f"Activation {spec.name!r} has no closed-form derivative kernel"
        )
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    u = scale * (z - center)
    return weight * (scale ** order) * spec.fastpath(u, order)

FUNCTION

refine

omnibias.torch.refine.refineSource ↗

refine(bank: AdaptivePackBank, residual_fn: ResidualFn, probes: Tensor, policy: RefinePolicy, *, step: int, target_fn: ResidualFn | None=None, probe_jet: tuple[float, Sequence[float]] | None=None) -> RefineReport

One refine tick. step is recorded for callers; ages increment here.

Singularity / scale-flow read target_fn when given (the field being resolved), otherwise the residual jet. Residual / certified-error read residual_fn. probe_jet overrides the jet (used by G3).

Parameters

Name	Type & default	Contract
bankpositional-or-keyword	AdaptivePackBankRequired	See the operation contract and implementation below.
residual_fnpositional-or-keyword	ResidualFnRequired	See the operation contract and implementation below.
probespositional-or-keyword	TensorRequired	See the operation contract and implementation below.
policypositional-or-keyword	RefinePolicyRequired	See the operation contract and implementation below.
stepkeyword-only	intRequired	See the operation contract and implementation below.
target_fnkeyword-only	ResidualFn | NoneDefault: None	See the operation contract and implementation below.
probe_jetkeyword-only	tuple[float, Sequence[float]] | NoneDefault: None	See the operation contract and implementation below.

ReturnsRefineReport

Explicit validation & failure cases

• When residual.shape != probes_1d.shape, raises ValueError(f'residual shape {tuple(residual.shape)} != probes {tuple(probes_1d.shape)}').

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

Executed regression · residual indicator misses scale

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_residual_indicator_misses_scale() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 65)

    def residual(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(-100.0 * x)

    report = refine(
        bank,
        residual,
        probes,
        RefinePolicy(
            indicator=Indicator.RESIDUAL,
            birth_threshold=0.1,
            death_threshold=0.0,
            min_age=10_000,
            hysteresis=1.0,
            min_scale_ratio=1.0,
        ),
        step=0,
    )
    assert report.proposed_scale is not None
    assert abs(report.proposed_scale) / 100.0 < 0.5

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_residual_indicator_misses_scale" -q

Regression example · budget stability with death

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_budget_stability_with_death() -> None:
    bank = _initial_bank()
    probes = torch.linspace(0.0, 1.0, 17)
    policy = RefinePolicy(
        indicator=Indicator.RESIDUAL,
        birth_threshold=0.5,
        death_threshold=0.2,
        min_age=2,
        hysteresis=1.0,
        max_packs=6,
        min_scale_ratio=1.0,
    )
    counts: list[int] = []
    for step in range(24):
        amp = 0.8 if step % 4 == 0 else 0.0

        def residual(x: torch.Tensor, _amp: float = amp) -> torch.Tensor:
            return torch.full_like(x, _amp)

        refine(bank, residual, probes, policy, step=step)
        counts.append(bank.n_active)
    assert max(counts) <= 6
    assert counts[-1] <= counts[0] + 2

python -m pytest packages/omnibias-torch/tests/test_refine.py -k test_budget_stability_with_death -q

Read the implementation

def refine(
    bank: AdaptivePackBank,
    residual_fn: ResidualFn,
    probes: Tensor,
    policy: RefinePolicy,
    *,
    step: int,
    target_fn: ResidualFn | None = None,
    probe_jet: tuple[float, Sequence[float]] | None = None,
) -> RefineReport:
    """One refine tick. ``step`` is recorded for callers; ages increment here.

    Singularity / scale-flow read ``target_fn`` when given (the field being
    resolved), otherwise the residual jet. Residual / certified-error read
    ``residual_fn``. ``probe_jet`` overrides the jet (used by G3).
    """
    del step
    probes_1d = _as_1d(probes)
    residual = _as_1d(residual_fn(probes_1d))
    if residual.shape != probes_1d.shape:
        raise ValueError(
            f"residual shape {tuple(residual.shape)} != probes {tuple(probes_1d.shape)}"
        )
    before = None
    if policy.debug_zero_perturbation:
        before = [float(v) for v in bank(probes_1d).detach().reshape(-1)]

    packs = bank.active_packs()
    field = _as_1d(bank(probes_1d)).detach()
    field_norm = float(field.square().sum().sqrt().item())
    contrib = pack_contributions(bank.pack_term_norms(probes_1d), field_norm)
    residuals = [float(v) for v in residual.detach()]
    probe_list = [float(v) for v in probes_1d.detach()]

    certified: list[float] | None = None
    if policy.indicator is Indicator.CERTIFIED_ERROR:
        if len(probe_list) < 2:
            spacing = 1.0
        else:
            spacing = abs(probe_list[1] - probe_list[0])
        certified = []
        for x in probe_list:
            jet = _jet_source(residual_fn, None, x, policy.jet_order)
            certified.append(certified_local_error(jet, spacing))

    peak_location: float | None = None
    derivatives: list[float] | None = None
    if probe_jet is not None:
        peak_location, derivatives = float(probe_jet[0]), [float(v) for v in probe_jet[1]]
    elif policy.indicator in (Indicator.SINGULARITY, Indicator.SCALE_FLOW):
        values = residuals if certified is None else certified
        peak_i = max(range(len(values)), key=lambda i: abs(values[i]))
        peak_location = probe_list[peak_i]
        derivatives = _jet_source(residual_fn, target_fn, peak_location, policy.jet_order)

    working, report = refine_bank(
        packs,
        policy,
        probes=probe_list,
        residuals=residuals,
        contributions=contrib,
        derivatives=derivatives,
        certified_values=certified,
        peak_location=peak_location,
    )
    bank.sync_from(working)
    if before is not None and (report.born or report.grown) and not report.died:
        after = [float(v) for v in bank(probes_1d).detach().reshape(-1)]
        assert_zero_perturbation(before, after)
    return report

Related APIsomnibias.core.assert_zero_perturbation↗omnibias.core.refine.certified_local_error↗omnibias.core.refine.pack_contributions↗omnibias.core.refine_bank↗

FUNCTION

scalar_jet_1d

omnibias.torch.refine.scalar_jet_1dSource ↗

scalar_jet_1d(fn: ScalarFn, x0: float, order: int) -> list[float]

phi^(k)(x0) for k = 0 .. order via nested reverse-mode.

Parameters

Name	Type & default	Contract
fnpositional-or-keyword	ScalarFnRequired	See the operation contract and implementation below.
x0positional-or-keyword	floatRequired	See the operation contract and implementation below.
orderpositional-or-keyword	intRequired	See the operation contract and implementation below.

Returnslist[float]

Explicit validation & failure cases

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

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

Executed regression · scalar jet exp

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_scalar_jet_exp() -> None:
    def fn(x: torch.Tensor) -> torch.Tensor:
        return torch.exp(3.0 * x)

    derivs = scalar_jet_1d(fn, 0.0, 4)
    for k, val in enumerate(derivs):
        assert val == pytest.approx(3.0**k, rel=1e-10)

python -m pytest "packages/omnibias-torch/tests/test_refine.py::test_scalar_jet_exp" -q

Read the implementation

def scalar_jet_1d(fn: ScalarFn, x0: float, order: int) -> list[float]:
    """``phi^(k)(x0)`` for ``k = 0 .. order`` via nested reverse-mode."""
    if order < 0:
        raise ValueError(f"order must be >= 0, got {order}")
    x = torch.tensor(float(x0), dtype=torch.get_default_dtype())

    def restriction(s: Tensor) -> Tensor:
        out = fn(s.reshape(()))
        if out.ndim != 0:
            out = out.reshape(())
        return out

    derivs = [float(restriction(x).detach())]
    current: Callable[[Tensor], Tensor] = restriction
    for _ in range(order):
        current = grad(current)
        derivs.append(float(current(x)))
    return derivs

← Package overviewRunnable learning examples →