omnibias.binary.schedule

0.1.0a2AlphaPython >=3.10Apache-2.0

Binary and multilevel quantization with explicit hard-forward and surrogate-gradient contracts.

import omnibias.binary.schedule

Module contract & mathematical conventions

Beta-annealing schedules for the tanh quantization surrogates (pure-Python).

The hard quantizer *forward* is independent of beta; only the smooth tanh(beta z) *backward* surrogate depends on it. Annealing beta upward over training realises the beta -> inf homotopy: the surrogate gradient beta * (1 - tanh(beta z)^2) concentrates ever more tightly at the decision boundary (its integral is fixed, so it approaches the 2 * delta Dirac limit), interpolating from an easy, well-conditioned early objective to the exact straight-through limit. This is the same saturation captured by the ActivationSpec limit metadata (the lim operator of Phase 1).

Pure-Python and backend-agnostic, so the identical scheduler drives torch and jax training loops; it mirrors the step() -> scalar contract of omnibias.torch.training.k_scheduler.KGrowthScheduler.

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

BetaAnnealScheduler

omnibias.binary.BetaAnnealSchedulerSource ↗

class BetaAnnealScheduler

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.

Anneal the surrogate sharpness beta from beta_start to beta_end.

Parameters

beta_start, beta_end : float Endpoint sharpness values; both must be > 0. beta_end > beta_start sharpens (the usual soft-to-hard curriculum), but either ordering works. num_steps : int Number of steps over which to anneal. value(step) is clamped to the endpoints outside [0, num_steps]. schedule : {"linear", "exp", "cosine"} Interpolation law. "linear" is linear in beta; "exp" is geometric (linear in log beta, natural for a multiplicative sharpness); "cosine" eases in and out.

Parameters

Name	Type & default	Contract
`beta_start`dataclass field	`float`Required	Endpoint sharpness values; both must be `> 0`. `beta_end > beta_start` sharpens (the usual soft-to-hard curriculum), but either ordering works.
`beta_end`dataclass field	`float`Required	See the operation contract and implementation below.
`num_steps`dataclass field	`int`Required	int Number of steps over which to anneal. `value(step)` is clamped to the endpoints outside `[0, num_steps]`.
`schedule`dataclass field	`str`Default: 'linear'	{"linear", "exp", "cosine"} Interpolation law. `"linear"` is linear in `beta`; `"exp"` is geometric (linear in `log beta`, natural for a multiplicative sharpness); `"cosine"` eases in and out.

Executed regression · invalid args raise

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_invalid_args_raise(kwargs: dict, match: str) -> None:
    with pytest.raises(ValueError, match=match):
        BetaAnnealScheduler(**kwargs)

python -m pytest "packages/omnibias-binary/tests/test_schedule.py::test_invalid_args_raise[kwargs0-must be > 0]" -q

Regression example · top level and module export are same object

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_top_level_and_module_export_are_same_object() -> None:
    assert BetaAnnealScheduler is DirectImport

python -m pytest packages/omnibias-binary/tests/test_schedule.py -k test_top_level_and_module_export_are_same_object -q

Fields & methods

beta_startfield

beta_start: float

Endpoint sharpness values; both must be > 0. beta_end > beta_start sharpens (the usual soft-to-hard curriculum), but either ordering works.

View source ↗

beta_endfield

beta_end: float

Endpoint sharpness values; both must be > 0. beta_end > beta_start sharpens (the usual soft-to-hard curriculum), but either ordering works.

View source ↗

num_stepsfield

num_steps: int

Number of steps over which to anneal. value(step) is clamped to the endpoints outside [0, num_steps].

View source ↗

schedulefield

schedule: str = 'linear'

Interpolation law. "linear" is linear in beta; "exp" is geometric (linear in log beta, natural for a multiplicative sharpness); "cosine" eases in and out.

View source ↗

__post_init__method

__post_init__(self) -> None

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

Dataclass construction calls this validation hook automatically. It rejects nonpositive endpoint beta values, fewer than one interpolation step, and schedule names outside linear, exp, and cosine. Construct the scheduler with valid finite endpoints rather than calling this method as a training-loop step; it neither computes a beta value nor advances the counter.

Returns: None

Validation & failure cases

When self.beta_start <= 0.0 or self.beta_end <= 0.0, raises ValueError('beta_start and beta_end must be > 0').
When self.num_steps < 1, raises ValueError(f'num_steps must be >= 1, got {self.num_steps}').
When self.schedule not in _SCHEDULES, raises ValueError(f'unknown schedule {self.schedule!r}; choose from {_SCHEDULES}').

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓

"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")

python binary_torch.py

Method implementation

    def __post_init__(self) -> None:
        if self.beta_start <= 0.0 or self.beta_end <= 0.0:
            raise ValueError("beta_start and beta_end must be > 0")
        if self.num_steps < 1:
            raise ValueError(f"num_steps must be >= 1, got {self.num_steps}")
        if self.schedule not in _SCHEDULES:
            raise ValueError(f"unknown schedule {self.schedule!r}; choose from {_SCHEDULES}")

View source ↗

valuemethod

value(self, step: int) -> float

Beta at an arbitrary step (clamped to the endpoints; stateless).

Read the schedule at an arbitrary numeric step without changing current_step. The normalized fraction is clamped to [0, 1], so negative steps return beta_start and steps beyond num_steps return beta_end. Linear interpolation acts on beta, exponential interpolation on log(beta), and cosine interpolation eases between endpoints. This method returns a Python float and does not construct a backend gradient graph.

Parameters

Name	Type & default	Contract
`step`positional-or-keyword	`int`Required	Position in the schedule, clamped to the endpoint interval without changing the internal call counter.

Returns: float

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓

"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")

python binary_torch.py

Method implementation

    def value(self, step: int) -> float:
        """Beta at an arbitrary step (clamped to the endpoints; stateless)."""
        frac = min(max(step / self.num_steps, 0.0), 1.0)
        if self.schedule == "linear":
            return self.beta_start + frac * (self.beta_end - self.beta_start)
        if self.schedule == "exp":
            log_beta = math.log(self.beta_start) + frac * (
                math.log(self.beta_end) - math.log(self.beta_start)
            )
            return math.exp(log_beta)
        # cosine ease-in/out: monotone 0 -> 1 as frac 0 -> 1.
        ease = 0.5 * (1.0 - math.cos(math.pi * frac))
        return self.beta_start + ease * (self.beta_end - self.beta_start)

View source ↗

stepmethod

step(self) -> float

Return the current beta and advance the internal step counter.

Obtain the beta value at the current counter, then increment that counter by one. The first call evaluates step zero, and calls after num_steps continue returning beta_end. Invoke once per chosen training unit, such as an optimizer update; the scheduler does not infer epochs or batches. The returned Python scalar controls the surrogate, while the hard quantizer forward remains unchanged.

Returns: float

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓

"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")

python binary_torch.py

Method implementation

    def step(self) -> float:
        """Return the current beta and advance the internal step counter."""
        beta = self.value(self._step)
        self._step += 1
        return beta

View source ↗

current_stepproperty

current_step(self) -> int

Number of times step has been called.

Read the number of calls made to step since construction or the most recent reset. This read-only integer property counts calls, not completed optimizer updates, and can exceed num_steps after the schedule has reached its endpoint. Use it when logging or coordinating the schedule; reading it does not advance interpolation.

Returns: int

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓

"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")

python binary_torch.py

Method implementation

    def current_step(self) -> int:
        """Number of times :meth:`step` has been called."""
        return self._step

View source ↗

resetmethod

reset(self) -> None

Rewind the internal step counter to zero.

Restart the schedule’s internal counter at zero. Endpoint values and the selected interpolation law are retained, and the next call to step returns beta_start. Use this when beginning a new training run with the same configuration. This mutating method returns None and does not reset model parameters or optimizer state.

Returns: None

Complete runnable example · binary torch

This is the complete downloadable script, tested against the pinned installed distribution. It includes initialization and assertions.

Download complete script ↓

"""Hard quantization, analytic surrogate derivatives, and an annealing schedule.

python -m pip install --pre "omnibias-binary==0.1.0a2" "omnibias-torch==0.5.0rc1"
"""
import math
import torch
from omnibias.binary import BetaAnnealScheduler
from omnibias.binary.torch.ops import (
    binarize, binarize01, heaviside, ternarize, kbit_quantize,
    riccati_tanh_derivative, riccati_sigmoid_derivative,
    surrogate_tower, surrogate_jet, curvature_corrected_slope, binarize_curvature,
)

z = torch.tensor([-0.75, -0.25, 0.0, 0.25, 0.75], dtype=torch.float64, requires_grad=True)
beta = 2.0
hard = binarize(z, beta)
torch.testing.assert_close(hard, torch.tensor([-1., -1., 1., 1., 1.], dtype=z.dtype))
# This is a specified surrogate gradient, not the classical derivative of sign.
(grad,) = torch.autograd.grad(hard.sum(), z)
torch.testing.assert_close(grad, beta * (1 - torch.tanh(beta*z).square()))
step = heaviside(z, beta)
torch.testing.assert_close(step, binarize01(z, beta))
(step_grad,) = torch.autograd.grad(step.sum(), z)
s = torch.sigmoid(beta*z)
torch.testing.assert_close(step_grad, beta*s*(1-s))
torch.testing.assert_close(ternarize(z, beta, delta=0.5), torch.tensor([-1., 0., 0., 0., 1.], dtype=z.dtype))
quantized = kbit_quantize(z, bits=2, lo=-1.0, hi=1.0, beta=beta)
assert quantized.shape == z.shape and torch.isfinite(quantized).all()
assert quantized.min() >= -1 and quantized.max() <= 1
(quantized_grad,) = torch.autograd.grad(quantized.sum(), z)
assert torch.isfinite(quantized_grad).all() and quantized_grad.abs().sum() > 0

t = torch.tanh(z)
torch.testing.assert_close(riccati_tanh_derivative(t, 1), 1-t.square())
s = torch.sigmoid(z)
torch.testing.assert_close(riccati_sigmoid_derivative(s, 1), s*(1-s))
tower = surrogate_tower(z, beta, order=3)
jet = surrogate_jet(z, beta, order=3)
assert tower.shape == jet.shape == (4, 5)
for order in range(4):
    torch.testing.assert_close(jet[order]*math.factorial(order), tower[order])
torch.testing.assert_close(curvature_corrected_slope(z, beta, window=0.0), tower[1])
(curvature_grad,) = torch.autograd.grad(binarize_curvature(z, beta).sum(), z)
torch.testing.assert_close(curvature_grad, curvature_corrected_slope(z, beta))

schedule = BetaAnnealScheduler(1.0, 4.0, num_steps=2, schedule="linear")
assert [schedule.step() for _ in range(3)] == [1.0, 2.5, 4.0]
assert schedule.current_step == 3 and schedule.value(100) == 4.0
schedule.reset()
assert schedule.current_step == 0 and schedule.step() == 1.0
print("quantizers: (5,); surrogate tower/jet: (4, 5); gradients and schedule passed")

python binary_torch.py

Method implementation

    def reset(self) -> None:
        """Rewind the internal step counter to zero."""
        self._step = 0

View source ↗

Read the implementation

class BetaAnnealScheduler:
    """Anneal the surrogate sharpness ``beta`` from ``beta_start`` to ``beta_end``.

    Parameters
    ----------
    beta_start, beta_end : float
        Endpoint sharpness values; both must be ``> 0``. ``beta_end > beta_start``
        sharpens (the usual soft-to-hard curriculum), but either ordering works.
    num_steps : int
        Number of steps over which to anneal. ``value(step)`` is clamped to the
        endpoints outside ``[0, num_steps]``.
    schedule : {"linear", "exp", "cosine"}
        Interpolation law. ``"linear"`` is linear in ``beta``; ``"exp"`` is
        geometric (linear in ``log beta``, natural for a multiplicative sharpness);
        ``"cosine"`` eases in and out.
    """

    beta_start: float
    beta_end: float
    num_steps: int
    schedule: str = "linear"
    _step: int = field(default=0, init=False, repr=False)

    def __post_init__(self) -> None:
        if self.beta_start <= 0.0 or self.beta_end <= 0.0:
            raise ValueError("beta_start and beta_end must be > 0")
        if self.num_steps < 1:
            raise ValueError(f"num_steps must be >= 1, got {self.num_steps}")
        if self.schedule not in _SCHEDULES:
            raise ValueError(f"unknown schedule {self.schedule!r}; choose from {_SCHEDULES}")

    def value(self, step: int) -> float:
        """Beta at an arbitrary step (clamped to the endpoints; stateless)."""
        frac = min(max(step / self.num_steps, 0.0), 1.0)
        if self.schedule == "linear":
            return self.beta_start + frac * (self.beta_end - self.beta_start)
        if self.schedule == "exp":
            log_beta = math.log(self.beta_start) + frac * (
                math.log(self.beta_end) - math.log(self.beta_start)
            )
            return math.exp(log_beta)
        # cosine ease-in/out: monotone 0 -> 1 as frac 0 -> 1.
        ease = 0.5 * (1.0 - math.cos(math.pi * frac))
        return self.beta_start + ease * (self.beta_end - self.beta_start)

    def step(self) -> float:
        """Return the current beta and advance the internal step counter."""
        beta = self.value(self._step)
        self._step += 1
        return beta

    @property
    def current_step(self) -> int:
        """Number of times :meth:`step` has been called."""
        return self._step

    def reset(self) -> None:
        """Rewind the internal step counter to zero."""
        self._step = 0

Other public import paths

omnibias.binary.BetaAnnealScheduler
omnibias.binary.schedule.BetaAnnealScheduler
omnibias.boolean.jax.ops.BetaAnnealScheduler
omnibias.boolean.torch.ops.BetaAnnealScheduler
omnibias.boolean.jax.ops.solver.BetaAnnealScheduler
omnibias.boolean.torch.ops.solver.BetaAnnealScheduler

← Package overview Runnable learning examples →

omnibias.binary.schedule

BetaAnnealScheduler​

Parameters

Fields & methods

Parameters

BetaAnnealScheduler