API REFERENCE/omnibias-keras

omnibias.keras.training.k_scheduler

0.0.2a1AlphaPython >=3.10Apache-2.0

Keras 3 operator layers and activation paths implemented with keras.ops for TensorFlow, JAX, or Torch.

import omnibias.keras.training.k_scheduler

Module contract & mathematical conventions

Plateau-triggered K-growth scheduler (Keras backend).

Watches a scalar metric; when it has not improved by min_delta for patience consecutive epochs, calls grow(strategy) on every GrowableOperatorMultiBiasUnit reachable from the model. After growth, an LR-boost window is opened. Mirrors omnibias.torch.training.k_scheduler.KGrowthScheduler; the training loop is responsible for reading current_lr_multiplier.

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 →

CONSTANT

AnchorProvider

omnibias.keras.training.AnchorProviderSource ↗

AnchorProvider = Callable[[str, GrowableOperatorMultiBiasUnit], 'float | np.ndarray | None']

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

AnchorProvider = Callable[
    [str, GrowableOperatorMultiBiasUnit], "float | np.ndarray | None"
]

Other public import paths

• omnibias.keras.training.AnchorProvider
• omnibias.keras.training.k_scheduler.AnchorProvider

CLASS

GrowthEvent

omnibias.keras.training.GrowthEventSource ↗

class GrowthEvent

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.

One growth event recorded by the scheduler.

Parameters

Name	Type & default	Contract
epochdataclass field	intRequired	See the operation contract and implementation below.
module_namedataclass field	strRequired	See the operation contract and implementation below.
new_active_Kdataclass field	intRequired	See the operation contract and implementation below.
strategydataclass field	strRequired	See the operation contract and implementation below.
anchordataclass field	float | NoneDefault: None	See the operation contract and implementation below.

Fields & methods

epochfield

epoch: int

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

View source ↗

module_namefield

module_name: str

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

View source ↗

new_active_Kfield

new_active_K: int

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

View source ↗

strategyfield

strategy: str

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

View source ↗

anchorfield

anchor: float | None = None

Declared float | None field on GrowthEvent. Default expression: None.

View source ↗

Read the implementation

class GrowthEvent:
    """One growth event recorded by the scheduler."""

    epoch: int
    module_name: str
    new_active_K: int
    strategy: str
    anchor: float | None = None

Other public import paths

• omnibias.keras.training.GrowthEvent
• omnibias.keras.training.k_scheduler.GrowthEvent

CLASS

KGrowthScheduler

omnibias.keras.KGrowthSchedulerSource ↗

class KGrowthScheduler

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.

Plateau-triggered K-growth controller (Keras backend).

Parameters mirror the torch scheduler: patience, min_delta, max_K, strategy, lr_boost_factor, lr_boost_epochs, cooldown_epochs, and an optional anchor_provider.

Parameters

Name	Type & default	Contract
modeldataclass field	layers.LayerRequired	See the operation contract and implementation below.
patiencedataclass field	intDefault: 10	See the operation contract and implementation below.
min_deltadataclass field	floatDefault: 0.001	See the operation contract and implementation below.
max_Kdataclass field	intDefault: 8	See the operation contract and implementation below.
strategydataclass field	GrowStrategyDefault: 'pair'	See the operation contract and implementation below.
lr_boost_factordataclass field	floatDefault: 10.0	See the operation contract and implementation below.
lr_boost_epochsdataclass field	intDefault: 5	See the operation contract and implementation below.
cooldown_epochsdataclass field	intDefault: 5	See the operation contract and implementation below.
anchor_providerdataclass field	AnchorProvider | NoneDefault: None	See the operation contract and implementation below.
best_lossdataclass field	floatDefault: float('inf')	See the operation contract and implementation below.
epochs_since_improvedataclass field	intDefault: 0	See the operation contract and implementation below.
cooldown_remainingdataclass field	intDefault: 0	See the operation contract and implementation below.
boost_remainingdataclass field	intDefault: 0	See the operation contract and implementation below.
growth_historydataclass field	list[GrowthEvent]Default: field(default_factory=list)	See the operation contract and implementation below.

Executed regression · scheduler requires growable unit

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_scheduler_requires_growable_unit() -> None:
    model = keras.Sequential(
        [keras.layers.Input(shape=(4,)), keras.layers.Dense(4)]
    )
    with pytest.raises(ValueError):
        KGrowthScheduler(model)

python -m pytest "packages/omnibias-keras/tests/test_growable.py::test_scheduler_requires_growable_unit" -q

Regression example · scheduler grows on plateau

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_scheduler_grows_on_plateau() -> None:
    model = keras.Sequential(
        [keras.layers.Input(shape=(4,)), GrowableOMBU(num_channels=4, init_K=1, K_max=5)]
    )
    sched = KGrowthScheduler(model, patience=2, min_delta=1e-3, max_K=5, strategy="pair")
    # Feed a flat metric to trigger a plateau.
    for epoch in range(6):
        sched.step(1.0, epoch=epoch)
    assert sched.total_growth_events() >= 1

python -m pytest packages/omnibias-keras/tests/test_growable.py -k test_scheduler_grows_on_plateau -q

Fields & methods

modelfield

model: layers.Layer

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

View source ↗

patiencefield

patience: int = 10

Declared int field on KGrowthScheduler. Default expression: 10.

View source ↗

min_deltafield

min_delta: float = 0.001

Declared float field on KGrowthScheduler. Default expression: 0.001.

View source ↗

max_Kfield

max_K: int = 8

Declared int field on KGrowthScheduler. Default expression: 8.

View source ↗

strategyfield

strategy: GrowStrategy = 'pair'

Declared GrowStrategy field on KGrowthScheduler. Default expression: 'pair'.

View source ↗

lr_boost_factorfield

lr_boost_factor: float = 10.0

Declared float field on KGrowthScheduler. Default expression: 10.0.

View source ↗

lr_boost_epochsfield

lr_boost_epochs: int = 5

Declared int field on KGrowthScheduler. Default expression: 5.

View source ↗

cooldown_epochsfield

cooldown_epochs: int = 5

Declared int field on KGrowthScheduler. Default expression: 5.

View source ↗

anchor_providerfield

anchor_provider: AnchorProvider | None = None

Declared AnchorProvider | None field on KGrowthScheduler. Default expression: None.

View source ↗

best_lossfield

best_loss: float = float('inf')

Declared float field on KGrowthScheduler. Default expression: float('inf').

View source ↗

epochs_since_improvefield

epochs_since_improve: int = 0

Declared int field on KGrowthScheduler. Default expression: 0.

View source ↗

cooldown_remainingfield

cooldown_remaining: int = 0

Declared int field on KGrowthScheduler. Default expression: 0.

View source ↗

boost_remainingfield

boost_remaining: int = 0

Declared int field on KGrowthScheduler. Default expression: 0.

View source ↗

growth_historyfield

growth_history: list[GrowthEvent] = field(default_factory=list)

Declared list[GrowthEvent] field on KGrowthScheduler. Default expression: field(default_factory=list).

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 not self._growable_units(), raises ValueError('KGrowthScheduler: the model contains no GrowableOperatorMultiBiasUnit children; nothing to grow.').

Executed regression · scheduler requires growable unit

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_scheduler_requires_growable_unit() -> None:
    model = keras.Sequential(
        [keras.layers.Input(shape=(4,)), keras.layers.Dense(4)]
    )
    with pytest.raises(ValueError):
        KGrowthScheduler(model)

python -m pytest "packages/omnibias-keras/tests/test_growable.py::test_scheduler_requires_growable_unit" -q

Method implementation

    def __post_init__(self) -> None:
        if not self._growable_units():
            raise ValueError(
                "KGrowthScheduler: the model contains no "
                "GrowableOperatorMultiBiasUnit children; nothing to grow."
            )

View source ↗

current_lr_multiplierproperty

current_lr_multiplier(self) -> float

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

Returns: float

Executed regression · scheduler grows on plateau

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_scheduler_grows_on_plateau() -> None:
    model = keras.Sequential(
        [keras.layers.Input(shape=(4,)), GrowableOMBU(num_channels=4, init_K=1, K_max=5)]
    )
    sched = KGrowthScheduler(model, patience=2, min_delta=1e-3, max_K=5, strategy="pair")
    # Feed a flat metric to trigger a plateau.
    for epoch in range(6):
        sched.step(1.0, epoch=epoch)
    assert sched.total_growth_events() >= 1

python -m pytest "packages/omnibias-keras/tests/test_growable.py::test_scheduler_grows_on_plateau" -q

Method implementation

    def current_lr_multiplier(self) -> float:
        return self.lr_boost_factor if self.boost_remaining > 0 else 1.0

View source ↗

stepmethod

step(self, metric: float, *, epoch: int) -> float

Update with the latest metric; return the LR multiplier to apply next.

Parameters

Name	Type & default	Contract
metricpositional-or-keyword	floatRequired	See the operation contract and implementation below.
epochkeyword-only	intRequired	See the operation contract and implementation below.

Returns: float

Executed regression · scheduler grows on plateau

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_scheduler_grows_on_plateau() -> None:
    model = keras.Sequential(
        [keras.layers.Input(shape=(4,)), GrowableOMBU(num_channels=4, init_K=1, K_max=5)]
    )
    sched = KGrowthScheduler(model, patience=2, min_delta=1e-3, max_K=5, strategy="pair")
    # Feed a flat metric to trigger a plateau.
    for epoch in range(6):
        sched.step(1.0, epoch=epoch)
    assert sched.total_growth_events() >= 1

python -m pytest "packages/omnibias-keras/tests/test_growable.py::test_scheduler_grows_on_plateau" -q

Method implementation

    def step(self, metric: float, *, epoch: int) -> float:
        """Update with the latest metric; return the LR multiplier to apply next."""
        if self.boost_remaining > 0:
            self.boost_remaining -= 1

        if not self._initialized:
            self.best_loss = metric
            self.epochs_since_improve = 0
            self._initialized = True
        else:
            improved = metric < self.best_loss * (1.0 - self.min_delta)
            if improved:
                self.best_loss = metric
                self.epochs_since_improve = 0
            else:
                self.epochs_since_improve += 1

        ready = (
            self.epochs_since_improve >= self.patience
            and self.cooldown_remaining == 0
            and self._has_capacity()
        )
        grew_this_step = False
        if ready:
            grew_this_step = self._grow_all_eligible(epoch=epoch)
            if grew_this_step:
                self.epochs_since_improve = 0
                self.boost_remaining = self.lr_boost_epochs
                self.cooldown_remaining = self.cooldown_epochs

        if not grew_this_step and self.cooldown_remaining > 0:
            self.cooldown_remaining -= 1

        return self.current_lr_multiplier

View source ↗

reset_bestmethod

reset_best(self) -> None

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

Returns: None

Method implementation

    def reset_best(self) -> None:
        self.best_loss = float("inf")
        self.epochs_since_improve = 0
        self._initialized = False

View source ↗

active_K_summarymethod

active_K_summary(self) -> dict[str, int]

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

Returns: dict[str, int]

Method implementation

    def active_K_summary(self) -> dict[str, int]:
        return {name: mod.active_K for name, mod in self._growable_units()}

View source ↗

total_growth_eventsmethod

total_growth_events(self) -> int

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

Returns: int

Executed regression · scheduler grows on plateau

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_scheduler_grows_on_plateau() -> None:
    model = keras.Sequential(
        [keras.layers.Input(shape=(4,)), GrowableOMBU(num_channels=4, init_K=1, K_max=5)]
    )
    sched = KGrowthScheduler(model, patience=2, min_delta=1e-3, max_K=5, strategy="pair")
    # Feed a flat metric to trigger a plateau.
    for epoch in range(6):
        sched.step(1.0, epoch=epoch)
    assert sched.total_growth_events() >= 1

python -m pytest "packages/omnibias-keras/tests/test_growable.py::test_scheduler_grows_on_plateau" -q

Method implementation

    def total_growth_events(self) -> int:
        return len(self.growth_history)

View source ↗

Read the implementation

class KGrowthScheduler:
    """Plateau-triggered K-growth controller (Keras backend).

    Parameters mirror the torch scheduler: ``patience``, ``min_delta``,
    ``max_K``, ``strategy``, ``lr_boost_factor``, ``lr_boost_epochs``,
    ``cooldown_epochs``, and an optional ``anchor_provider``.
    """

    model: layers.Layer
    patience: int = 10
    min_delta: float = 1e-3
    max_K: int = 8
    strategy: GrowStrategy = "pair"
    lr_boost_factor: float = 10.0
    lr_boost_epochs: int = 5
    cooldown_epochs: int = 5
    anchor_provider: AnchorProvider | None = None

    best_loss: float = float("inf")
    epochs_since_improve: int = 0
    cooldown_remaining: int = 0
    boost_remaining: int = 0
    growth_history: list[GrowthEvent] = field(default_factory=list)
    _initialized: bool = False

    def __post_init__(self) -> None:
        if not self._growable_units():
            raise ValueError(
                "KGrowthScheduler: the model contains no "
                "GrowableOperatorMultiBiasUnit children; nothing to grow."
            )

    @property
    def current_lr_multiplier(self) -> float:
        return self.lr_boost_factor if self.boost_remaining > 0 else 1.0

    def step(self, metric: float, *, epoch: int) -> float:
        """Update with the latest metric; return the LR multiplier to apply next."""
        if self.boost_remaining > 0:
            self.boost_remaining -= 1

        if not self._initialized:
            self.best_loss = metric
            self.epochs_since_improve = 0
            self._initialized = True
        else:
            improved = metric < self.best_loss * (1.0 - self.min_delta)
            if improved:
                self.best_loss = metric
                self.epochs_since_improve = 0
            else:
                self.epochs_since_improve += 1

        ready = (
            self.epochs_since_improve >= self.patience
            and self.cooldown_remaining == 0
            and self._has_capacity()
        )
        grew_this_step = False
        if ready:
            grew_this_step = self._grow_all_eligible(epoch=epoch)
            if grew_this_step:
                self.epochs_since_improve = 0
                self.boost_remaining = self.lr_boost_epochs
                self.cooldown_remaining = self.cooldown_epochs

        if not grew_this_step and self.cooldown_remaining > 0:
            self.cooldown_remaining -= 1

        return self.current_lr_multiplier

    def reset_best(self) -> None:
        self.best_loss = float("inf")
        self.epochs_since_improve = 0
        self._initialized = False

    def active_K_summary(self) -> dict[str, int]:
        return {name: mod.active_K for name, mod in self._growable_units()}

    def total_growth_events(self) -> int:
        return len(self.growth_history)

    # ----- internals -------------------------------------------------------

    def _growable_units(self) -> list[tuple[str, GrowableOperatorMultiBiasUnit]]:
        return [
            (name, mod)
            for name, mod in _iter_sublayers(self.model)
            if isinstance(mod, GrowableOperatorMultiBiasUnit)
        ]

    def _strategy_increment(self) -> int:
        return 2 if self.strategy == "pair" else 1

    def _has_capacity(self) -> bool:
        inc = self._strategy_increment()
        for _, mod in self._growable_units():
            if mod.active_K + inc <= min(self.max_K, mod.K_max):
                return True
        return False

    def _grow_all_eligible(self, *, epoch: int) -> bool:
        inc = self._strategy_increment()
        grew_any = False
        for name, mod in self._growable_units():
            if mod.active_K + inc > min(self.max_K, mod.K_max):
                continue
            anchor: float | np.ndarray | None = None
            if self.anchor_provider is not None and self.strategy == "pair":
                try:
                    anchor = self.anchor_provider(name, mod)
                except Exception as exc:  # noqa: BLE001
                    # An anchor provider failure must not abort training; log and
                    # fall back to the default anchor (anchor=None).
                    _LOGGER.warning(
                        "anchor_provider failed for %r: %r; falling back to "
                        "default anchor.",
                        name,
                        exc,
                    )
                    anchor = None
            try:
                if anchor is not None:
                    mod.grow(strategy=self.strategy, anchor_value=anchor)
                else:
                    mod.grow(strategy=self.strategy)
            except (RuntimeError, ValueError):
                continue
            anchor_scalar: float | None = None
            if isinstance(anchor, np.ndarray):
                anchor_scalar = float(anchor.mean())
            elif anchor is not None:
                anchor_scalar = float(anchor)
            self.growth_history.append(
                GrowthEvent(
                    epoch=epoch,
                    module_name=name,
                    new_active_K=mod.active_K,
                    strategy=self.strategy,
                    anchor=anchor_scalar,
                )
            )
            grew_any = True
        return grew_any

Related APIsomnibias.keras.training.GrowthEvent↗

Other public import paths

• omnibias.keras.KGrowthScheduler
• omnibias.keras.training.KGrowthScheduler
• omnibias.keras.training.k_scheduler.KGrowthScheduler

← Package overviewRunnable learning examples →