lrsched 0.1.0__py3-none-any.whl

lrsched/__init__.py

@@ -0,0 +1,34 @@
+"""Framework-agnostic learning-rate schedules as pure functions, with zero dependencies."""
+
+from ._types import Schedule
+from .compose import sample, sequential, with_warmup
+from .schedules import (
+    constant,
+    cosine,
+    cosine_restarts,
+    exponential,
+    inverse_sqrt,
+    linear,
+    multi_step,
+    one_cycle,
+    polynomial,
+    step_decay,
+)
+
+__all__ = [
+    "Schedule",
+    "constant",
+    "cosine",
+    "cosine_restarts",
+    "exponential",
+    "inverse_sqrt",
+    "linear",
+    "multi_step",
+    "one_cycle",
+    "polynomial",
+    "sample",
+    "sequential",
+    "step_decay",
+    "with_warmup",
+]
+__version__ = "0.1.0"

lrsched/_types.py

@@ -0,0 +1,4 @@
+from collections.abc import Callable
+
+# A schedule maps a non-negative integer step to a learning rate.
+Schedule = Callable[[int], float]

lrsched/_validate.py

@@ -0,0 +1,19 @@
+from math import isfinite
+
+
+def check_finite(name: str, value: float) -> float:
+    if not isfinite(value):
+        raise ValueError(f"{name} must be a finite number, received {value!r}")
+    return value
+
+
+def check_positive_int(name: str, value: int) -> int:
+    if not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer, received {value!r}")
+    return value
+
+
+def check_step(step: int) -> int:
+    if not isinstance(step, int) or step < 0:
+        raise ValueError(f"step must be a non-negative integer, received {step!r}")
+    return step

lrsched/compose.py

@@ -0,0 +1,48 @@
+from collections.abc import Sequence
+
+from ._types import Schedule
+from ._validate import check_finite, check_positive_int, check_step
+
+
+def with_warmup(schedule: Schedule, *, warmup_steps: int, start_lr: float) -> Schedule:
+    """Prepend a linear warmup from start_lr to the wrapped schedule's value at step 0.
+    After warmup, the wrapped schedule runs with its step counted from the warmup end."""
+    check_positive_int("warmup_steps", warmup_steps)
+    check_finite("start_lr", start_lr)
+    target = schedule(0)
+
+    def wrapped(step: int) -> float:
+        check_step(step)
+        if step < warmup_steps:
+            return start_lr + (target - start_lr) * (step / warmup_steps)
+        return schedule(step - warmup_steps)
+
+    return wrapped
+
+
+def sequential(phases: Sequence[tuple[int, Schedule]]) -> Schedule:
+    """Run schedules back to back. Each phase is (duration, schedule), and within a phase
+    the schedule sees a step counted from that phase's start. Past the final phase, the
+    last phase's final value is held."""
+    if len(phases) == 0:
+        raise ValueError("sequential requires at least one phase")
+    for duration, _ in phases:
+        check_positive_int("phase duration", duration)
+
+    def wrapped(step: int) -> float:
+        check_step(step)
+        offset = 0
+        for duration, schedule in phases:
+            if step < offset + duration:
+                return schedule(step - offset)
+            offset += duration
+        last_duration, last_schedule = phases[-1]
+        return last_schedule(last_duration - 1)
+
+    return wrapped
+
+
+def sample(schedule: Schedule, *, num_steps: int) -> list[float]:
+    """Evaluate a schedule at steps 0..num_steps-1, useful for plotting or testing."""
+    check_positive_int("num_steps", num_steps)
+    return [schedule(i) for i in range(num_steps)]

lrsched/schedules.py

@@ -0,0 +1,170 @@
+from collections.abc import Sequence
+from math import cos, pi
+
+from ._types import Schedule
+from ._validate import check_finite, check_positive_int, check_step
+
+
+def constant(*, lr: float) -> Schedule:
+    """A flat learning rate."""
+    check_finite("lr", lr)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        return lr
+
+    return schedule
+
+
+def step_decay(*, base_lr: float, gamma: float, step_size: int) -> Schedule:
+    """Multiply by gamma every step_size steps."""
+    check_finite("base_lr", base_lr)
+    check_finite("gamma", gamma)
+    check_positive_int("step_size", step_size)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        return base_lr * gamma ** (step // step_size)
+
+    return schedule
+
+
+def multi_step(*, base_lr: float, milestones: Sequence[int], gamma: float) -> Schedule:
+    """Multiply by gamma at each milestone step."""
+    check_finite("base_lr", base_lr)
+    check_finite("gamma", gamma)
+    ordered = sorted(milestones)
+    for m in ordered:
+        check_positive_int("milestone", m)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        drops = sum(1 for m in ordered if step >= m)
+        return base_lr * gamma**drops
+
+    return schedule
+
+
+def exponential(*, base_lr: float, gamma: float) -> Schedule:
+    """Decay by a constant factor gamma each step: base_lr * gamma ** step."""
+    check_finite("base_lr", base_lr)
+    check_finite("gamma", gamma)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        return base_lr * gamma**step
+
+    return schedule
+
+
+def polynomial(*, base_lr: float, end_lr: float, total_steps: int, power: float) -> Schedule:
+    """Polynomial decay from base_lr to end_lr over total_steps, then hold end_lr."""
+    check_finite("base_lr", base_lr)
+    check_finite("end_lr", end_lr)
+    check_positive_int("total_steps", total_steps)
+    check_finite("power", power)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        t = min(step, total_steps)
+        remaining = 1.0 - t / total_steps
+        return float(end_lr + (base_lr - end_lr) * remaining**power)
+
+    return schedule
+
+
+def linear(*, base_lr: float, end_lr: float, total_steps: int) -> Schedule:
+    """Linear interpolation from base_lr to end_lr over total_steps, then hold end_lr."""
+    check_finite("base_lr", base_lr)
+    check_finite("end_lr", end_lr)
+    check_positive_int("total_steps", total_steps)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        t = min(step, total_steps)
+        return base_lr + (end_lr - base_lr) * (t / total_steps)
+
+    return schedule
+
+
+def cosine(*, base_lr: float, min_lr: float, total_steps: int) -> Schedule:
+    """Cosine annealing from base_lr at step 0 to min_lr at total_steps, then hold."""
+    check_finite("base_lr", base_lr)
+    check_finite("min_lr", min_lr)
+    check_positive_int("total_steps", total_steps)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        t = min(step, total_steps)
+        return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + cos(pi * t / total_steps))
+
+    return schedule
+
+
+def cosine_restarts(*, base_lr: float, min_lr: float, period: int, t_mult: int) -> Schedule:
+    """Cosine annealing with warm restarts (SGDR). Each cycle resets to base_lr; cycle
+    length is multiplied by t_mult after each restart (t_mult = 1 keeps it periodic)."""
+    check_finite("base_lr", base_lr)
+    check_finite("min_lr", min_lr)
+    check_positive_int("period", period)
+    check_positive_int("t_mult", t_mult)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        start = 0
+        current = period
+        while step >= start + current:
+            start += current
+            current *= t_mult
+        t = step - start
+        return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + cos(pi * t / current))
+
+    return schedule
+
+
+def inverse_sqrt(*, base_lr: float, warmup_steps: int) -> Schedule:
+    """Transformer schedule: linear warmup to base_lr at warmup_steps, then decay by the
+    inverse square root of the step."""
+    check_finite("base_lr", base_lr)
+    check_positive_int("warmup_steps", warmup_steps)
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        if step < warmup_steps:
+            return base_lr * step / warmup_steps
+        return float(base_lr * (warmup_steps / step) ** 0.5)
+
+    return schedule
+
+
+def one_cycle(
+    *, max_lr: float, total_steps: int, pct_start: float, initial_div: float, final_div: float
+) -> Schedule:
+    """One-cycle policy: cosine ramp up to max_lr over the first pct_start of training,
+    then cosine anneal down to max_lr / final_div. Starts at max_lr / initial_div."""
+    check_finite("max_lr", max_lr)
+    check_positive_int("total_steps", total_steps)
+    check_finite("pct_start", pct_start)
+    check_finite("initial_div", initial_div)
+    check_finite("final_div", final_div)
+    if not 0.0 < pct_start < 1.0:
+        raise ValueError(f"pct_start must be between 0 and 1, received {pct_start!r}")
+    if initial_div <= 0.0 or final_div <= 0.0:
+        raise ValueError("initial_div and final_div must be positive")
+    if total_steps < 2:
+        raise ValueError("total_steps must be at least 2 for one_cycle")
+
+    warmup = min(max(1, round(pct_start * total_steps)), total_steps - 1)
+    start_lr = max_lr / initial_div
+    end_lr = max_lr / final_div
+
+    def schedule(step: int) -> float:
+        check_step(step)
+        if step <= warmup:
+            frac = step / warmup
+            return start_lr + 0.5 * (max_lr - start_lr) * (1.0 - cos(pi * frac))
+        t = min(step, total_steps)
+        frac = (t - warmup) / (total_steps - warmup)
+        return end_lr + 0.5 * (max_lr - end_lr) * (1.0 + cos(pi * frac))
+
+    return schedule

lrsched-0.1.0.dist-info/METADATA

@@ -0,0 +1,146 @@
+Metadata-Version: 2.4
+Name: lrsched
+Version: 0.1.0
+Summary: Framework-agnostic learning-rate schedules as pure functions, in Python with zero dependencies.
+Project-URL: Homepage, https://github.com/amaar-mc/lrsched
+Project-URL: Repository, https://github.com/amaar-mc/lrsched
+Project-URL: Issues, https://github.com/amaar-mc/lrsched/issues
+Project-URL: Changelog, https://github.com/amaar-mc/lrsched/blob/main/CHANGELOG.md
+Author: Amaar Chughtai
+License: MIT License
+        
+        Copyright (c) 2026 Amaar Chughtai
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: cosine-annealing,deep-learning,learning-rate,lr-schedule,machine-learning,one-cycle,scheduler,sgdr,training,warmup
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Provides-Extra: dev
+Requires-Dist: hypothesis>=6; extra == 'dev'
+Requires-Dist: mypy>=1.11; extra == 'dev'
+Requires-Dist: pytest>=8; extra == 'dev'
+Requires-Dist: ruff>=0.6; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# lrsched
+
+<p align="center">
+  <img src="assets/logo.png" alt="lrsched logo" width="160">
+</p>
+
+[![PyPI](https://img.shields.io/pypi/v/lrsched)](https://pypi.org/project/lrsched/)
+[![CI](https://github.com/amaar-mc/lrsched/actions/workflows/ci.yml/badge.svg)](https://github.com/amaar-mc/lrsched/actions/workflows/ci.yml)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](./LICENSE)
+
+Framework-agnostic learning-rate schedules as pure functions, in Python with zero dependencies. Each schedule maps a step to a learning rate, so it works in any training loop or framework, or none.
+
+## Install
+
+```sh
+pip install lrsched
+```
+
+## 30-second example
+
+```python
+from lrsched import cosine, with_warmup, sample
+
+schedule = with_warmup(
+    cosine(base_lr=1e-3, min_lr=1e-5, total_steps=1000),
+    warmup_steps=100,
+    start_lr=0.0,
+)
+
+lr = schedule(250)            # learning rate at step 250
+curve = sample(schedule, num_steps=1000)  # the whole curve, for plotting or logging
+```
+
+A schedule is just `Callable[[int], float]`. Plug `schedule(step)` into your optimizer
+however your framework expects, or use it to drive a plain training loop.
+
+## Why this exists
+
+Every learning-rate scheduler is tied to a framework: `torch.optim.lr_scheduler`,
+`timm`, `transformers`, or `optax` for JAX. If you write a custom loop, use a
+non-PyTorch stack, or just want to plot a schedule, you end up pasting a `LambdaLR`
+snippet. `lrsched` is a small, dependency-free library where each schedule is a pure
+function, easy to test, plot, log, and reuse anywhere.
+
+## Comparison
+
+| | lrsched | torch / timm | optax |
+|---|:---:|:---:|:---:|
+| Framework | none | PyTorch | JAX |
+| Pure step to lr function | yes | no (optimizer-bound) | partial |
+| Zero dependencies | yes | no | no |
+| Composable warmup and phases | yes | partial | yes |
+
+## Schedules
+
+- `constant`, `step_decay`, `multi_step`, `exponential`
+- `linear`, `polynomial`
+- `cosine`, `cosine_restarts` (SGDR)
+- `inverse_sqrt` (Transformer)
+- `one_cycle`
+
+## Composition
+
+- `with_warmup(schedule, ...)` prepends a linear warmup.
+- `sequential(phases)` runs schedules back to back.
+- `sample(schedule, num_steps=...)` evaluates a schedule over a range.
+
+Parameters are required keyword arguments, so every schedule reads explicitly at the call
+site. Schedules hold their final value past the end rather than erroring, and a negative
+step raises.
+
+## Examples
+
+```sh
+python examples/schedules.py
+```
+
+## Testing
+
+```sh
+pip install -e ".[dev]"
+pytest
+```
+
+Tests cover the exact value of each schedule at known steps, schedule-specific shapes
+(restarts, the one-cycle peak, the warmup handoff), and invariants checked with
+Hypothesis (cosine stays within bounds, warmup is monotone).
+
+## Contributing
+
+Issues and pull requests are welcome. See [CONTRIBUTING.md](./CONTRIBUTING.md).
+
+## License
+
+MIT. See [LICENSE](./LICENSE).

lrsched-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+lrsched/__init__.py,sha256=IIce4y2wDGfX20N5mDrg4n_auf5XTz8Y0gE13yvORDs,640
+lrsched/_types.py,sha256=juJaHUF6QXBdcDR0x5PS9gkdqB9NhiVkmL999gNHCKA,138
+lrsched/_validate.py,sha256=5GC6kEL4Qx7BIM_tzbD-PDMJfOuuOQFvIBmblAw6sQA,592
+lrsched/compose.py,sha256=yhFPuT3FJQ5ierXeNAT_OGU-IfWAT-KjxU9K6rvXIhE,1835
+lrsched/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lrsched/schedules.py,sha256=gx5K_mc1Wu6HzAhjeJk_FAZiKJx1n21eWkuyZ1gRXKc,5761
+lrsched-0.1.0.dist-info/METADATA,sha256=oXungvasmPsK9ef61NjjOhEqaC6cFlaLouR2hlXMdxQ,5511
+lrsched-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+lrsched-0.1.0.dist-info/licenses/LICENSE,sha256=NORTVJb4x206RXQkpZt_vpj8I7aZL6Vn26BvTOq6J40,1071
+lrsched-0.1.0.dist-info/RECORD,,

lrsched-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

lrsched-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Amaar Chughtai
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.