lrsched 0.1.0__tar.gz

lrsched-0.1.0/.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,24 @@
+---
+name: Bug report
+about: A wrong learning rate or unexpected behavior
+title: ""
+labels: bug
+---
+
+## Input
+
+The schedule and parameters, for example
+`cosine(base_lr=1e-3, min_lr=1e-5, total_steps=1000)` at step 500.
+
+## Expected
+
+The learning rate you expected, with a source if it is a known formula.
+
+## Actual
+
+The value you observed.
+
+## Environment
+
+- lrsched version:
+- Python version:

lrsched-0.1.0/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,18 @@
+---
+name: Feature request
+about: Suggest a schedule or capability for lrsched
+title: ""
+labels: enhancement
+---
+
+## What
+
+The schedule or capability you would like.
+
+## Why
+
+The training or analysis workflow this would support.
+
+## References
+
+The paper or formula that pins down the expected behavior.

lrsched-0.1.0/.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,13 @@
+## Summary
+
+What this changes and why.
+
+## Checklist
+
+- [ ] Tests added or updated (a bug fix starts with a failing test)
+- [ ] Exact reference values for any new schedule, plus a shape check
+- [ ] `uv run ruff check .` passes
+- [ ] `uv run mypy src` passes
+- [ ] `uv run pytest` passes
+- [ ] No runtime dependencies added
+- [ ] No em dash characters in docs or commit messages

lrsched-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,22 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python: ["3.10", "3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+      - run: pip install -e ".[dev]"
+      - run: ruff check .
+      - run: mypy src
+      - run: pytest -q

lrsched-0.1.0/.gitignore

@@ -0,0 +1,15 @@
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+uv.lock
+build/
+dist/
+*.egg-info/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.hypothesis/
+.coverage
+htmlcov/
+.DS_Store

lrsched-0.1.0/CHANGELOG.md

@@ -0,0 +1,21 @@
+# Changelog
+
+All notable changes to this project are documented here. The format follows
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the project
+adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+### Planned
+- Cyclical triangular schedules.
+- A small CLI to print or sample a schedule.
+
+## [0.1.0]
+
+### Added
+- Schedule type (a callable from step to learning rate).
+- Schedules: constant, step_decay, multi_step, exponential, linear, polynomial, cosine,
+  cosine_restarts (SGDR), inverse_sqrt (Transformer), one_cycle.
+- Composition: with_warmup, sequential, sample.
+- Validation with clear errors. Test suite with exact reference values and Hypothesis
+  invariants.

lrsched-0.1.0/CLAUDE.md

@@ -0,0 +1,47 @@
+# lrsched
+
+Pure-Python, framework-agnostic learning-rate schedules. Zero runtime dependencies.
+
+## Commands
+
+- Create env and install: `uv venv && uv pip install -e ".[dev]"`
+- Test: `uv run pytest -q`
+- Lint: `uv run ruff check .` (format with `uv run ruff format .`)
+- Types: `uv run mypy src`
+- Build: `uv build` (then `uv run --with twine twine check dist/*` before publishing)
+
+## Architecture
+
+`src/lrsched/`:
+- `_types.py` the Schedule type (Callable[[int], float])
+- `_validate.py` shared validation (finite floats, positive ints, step)
+- `schedules.py` schedule factories (cosine, one_cycle, sgdr, inverse_sqrt, and so on)
+- `compose.py` with_warmup, sequential, sample
+- `__init__.py` public surface
+
+See `docs/architecture.md` for the formulas.
+
+## Conventions
+
+- Each schedule factory returns a pure function from step to learning rate.
+- Parameters are required keyword-only arguments; no default values.
+- Schedules hold their final value past the end; a negative step raises.
+- No runtime dependencies. Keep the public API small.
+
+## Testing rules
+
+- Exact value of each schedule at known steps.
+- Shape checks (restart boundaries, one-cycle peak, warmup handoff).
+- Hypothesis invariants (cosine within bounds, warmup monotone).
+- Bug fixes start with a failing test.
+
+## Release
+
+- Semantic versioning; update `CHANGELOG.md` and `__version__`.
+- Gates: `uv run pytest && uv run ruff check . && uv run mypy src && uv build && uv run twine check dist/*`.
+- Publish to PyPI, tag `vX.Y.Z`, GitHub release.
+
+## Style
+
+- No em dash characters in docs, comments, or commit messages.
+- Comments explain non-obvious reasoning only.

lrsched-0.1.0/CODE_OF_CONDUCT.md

@@ -0,0 +1,37 @@
+# Code of Conduct
+
+## Our pledge
+
+We as members, contributors, and maintainers pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity and
+orientation.
+
+## Our standards
+
+Examples of behavior that contributes to a positive environment:
+
+- Showing empathy and kindness toward other people.
+- Being respectful of differing opinions, viewpoints, and experiences.
+- Giving and gracefully accepting constructive feedback.
+- Focusing on what is best for the community.
+
+Examples of unacceptable behavior:
+
+- Harassment, insulting or derogatory comments, and personal or political attacks.
+- Public or private harassment.
+- Publishing others' private information without explicit permission.
+- Other conduct which could reasonably be considered inappropriate.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the maintainer at amaar2cool@gmail.com. All complaints will be
+reviewed and investigated promptly and fairly.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org),
+version 2.1.

lrsched-0.1.0/CONTRIBUTING.md

@@ -0,0 +1,37 @@
+# Contributing to lrsched
+
+Thanks for your interest. This project values correctness, a small surface area,
+and zero runtime dependencies.
+
+## Development
+
+```sh
+uv venv
+uv pip install -e ".[dev]"
+uv run pytest -q
+uv run ruff check .
+uv run mypy src
+```
+
+A standard virtual environment with `pip install -e ".[dev]"` works the same way.
+
+## Guidelines
+
+- No runtime dependencies.
+- Each schedule is a pure function from step to learning rate, with required
+  keyword-only parameters and no default values.
+- Every schedule needs an exact reference-value test at known steps, a shape check, and
+  any invariant it should satisfy.
+- A bug fix starts with a failing test.
+- Run `uv run ruff format .` before committing.
+- Commit messages follow `type(scope): description`.
+
+## Adding a schedule
+
+State the formula in the docstring and the README, give exact test values at a few steps,
+and document the behavior past the end of training.
+
+## Reporting issues
+
+Open an issue with the schedule and parameters, the step, the expected learning rate with
+a source, and the value you observed.

lrsched-0.1.0/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.

lrsched-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,95 @@
+# 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/SECURITY.md

@@ -0,0 +1,19 @@
+# Security Policy
+
+## Scope
+
+`lrsched` is a pure computation library with no runtime dependencies, no network
+access, and no file system access. The attack surface is limited to incorrect
+results from malformed input, which the library guards against with explicit
+validation.
+
+## Reporting a vulnerability
+
+If you find a security issue, please email amaar2cool@gmail.com with details and
+steps to reproduce. Please do not open a public issue for security reports. You
+can expect an initial response within a few days.
+
+## Supported versions
+
+The latest published minor version receives fixes. Pre-1.0 releases may introduce
+breaking changes in minor versions, as allowed by semantic versioning.

lrsched-0.1.0/docs/architecture.md

@@ -0,0 +1,50 @@
+# Architecture
+
+`lrsched` is a few small pure modules. The core idea is that a schedule is a function
+`Callable[[int], float]` from step to learning rate, built by a factory that captures its
+parameters in a closure.
+
+## Modules
+
+- `_types.py` defines `Schedule`.
+- `_validate.py` holds the shared checks (finite floats, positive integers, non-negative
+  step) so every factory reports bad input the same way.
+- `schedules.py` holds the factories. Each validates its parameters once, then returns a
+  closure that validates the step and computes the rate.
+- `compose.py` holds `with_warmup`, `sequential`, and `sample`, which operate on schedules
+  rather than parameters.
+
+## Formulas
+
+With base learning rate b, minimum m, total steps T, and step s clamped to T where noted:
+
+- step_decay: `b * gamma ** (s // step_size)`.
+- multi_step: `b * gamma ** (number of milestones <= s)`.
+- exponential: `b * gamma ** s`.
+- linear: `b + (end - b) * (min(s, T) / T)`.
+- polynomial: `end + (b - end) * (1 - min(s, T) / T) ** power`.
+- cosine: `m + 0.5 * (b - m) * (1 + cos(pi * min(s, T) / T))`.
+- cosine_restarts: the cosine formula within the current cycle, where cycle length starts
+  at `period` and is multiplied by `t_mult` after each restart. The current cycle is found
+  by stepping through cycle lengths until the step falls inside one.
+- inverse_sqrt: `b * s / warmup` during warmup, then `b * sqrt(warmup / s)`.
+- one_cycle: a cosine ramp from `max_lr / initial_div` up to `max_lr` over the first
+  `pct_start` of training, then a cosine anneal down to `max_lr / final_div`.
+
+## Composition
+
+`with_warmup` reads the wrapped schedule's value at step 0 as the warmup target, ramps
+linearly from `start_lr` to that target over the warmup, then runs the wrapped schedule
+with the step shifted so it begins at its own step 0 when warmup ends. `sequential` maps a
+global step into the active phase and passes a phase-local step to that phase's schedule.
+
+## Edges
+
+Finite schedules clamp the step to their total, so they hold their final value rather than
+extrapolating past the end. A negative step raises, since it is almost always a bug.
+
+## Why pure Python
+
+The formulas are small and exact. Implementing them with no dependencies means the same
+schedule runs under PyTorch, JAX, NumPy, or a plain loop, and tests are exact rather than
+approximate.

lrsched-0.1.0/docs/charter.md

@@ -0,0 +1,33 @@
+# Charter
+
+## Purpose
+
+Provide correct, dependency-free learning-rate schedules as pure functions, so that any
+training loop or framework, or none, can use the same schedule and so that schedules are
+easy to test, plot, and reason about.
+
+## Scope
+
+- The common schedules: constant, step, multi-step, exponential, linear, polynomial,
+  cosine, cosine with warm restarts, inverse square root, and one-cycle.
+- Composition: warmup, phase chaining, and sampling.
+
+## Non-goals
+
+- An optimizer or training framework. Schedules return a learning rate; the caller
+  applies it.
+- Framework coupling. There is no dependency on PyTorch, JAX, or any array library.
+- Stateful schedulers. Schedules are pure functions of the step.
+
+## Principles
+
+- Correctness first. Each schedule is tested against its exact formula at known steps.
+- Small, stable public API. Pure functions, explicit parameters.
+- Zero runtime dependencies.
+- Predictable edges. Schedules hold their final value past the end; a negative step
+  raises.
+
+## Audience
+
+People writing custom training loops, using non-PyTorch stacks, driving schedules from
+config, or teaching how learning-rate schedules behave.

lrsched-0.1.0/docs/logo-prompt.md

@@ -0,0 +1,10 @@
+# Logo prompt
+
+The logo in `assets/logo.png` was generated with OpenAI `gpt-image-1` (1024x1024,
+medium quality) from this prompt:
+
+> Minimal geometric logo for a learning-rate schedule library. A single smooth curve that
+> rises then falls, like a warmup followed by cosine decay, drawn over faint axes. Flat
+> vector, premium research-lab aesthetic. White background, black curve, one pale blue
+> accent dot at the peak. Subtle thin grid. Centered, generous negative space, no text, no
+> letters, no robots, no neon, crisp and high-signal.

lrsched-0.1.0/examples/schedules.py

@@ -0,0 +1,30 @@
+"""Sample a few schedules and print them as small ASCII sparklines.
+
+Run with: python examples/schedules.py
+"""
+
+from lrsched import cosine, one_cycle, sample, with_warmup
+
+BARS = " .:-=+*#%@"
+
+
+def sparkline(values: list[float]) -> str:
+    lo = min(values)
+    hi = max(values)
+    span = hi - lo or 1.0
+    return "".join(BARS[min(len(BARS) - 1, int((v - lo) / span * (len(BARS) - 1)))] for v in values)
+
+
+total = 60
+named = {
+    "cosine": cosine(base_lr=1.0, min_lr=0.0, total_steps=total),
+    "warmup+cosine": with_warmup(
+        cosine(base_lr=1.0, min_lr=0.0, total_steps=total - 10), warmup_steps=10, start_lr=0.0
+    ),
+    "one_cycle": one_cycle(
+        max_lr=1.0, total_steps=total, pct_start=0.3, initial_div=25.0, final_div=1000.0
+    ),
+}
+
+for name, schedule in named.items():
+    print(f"{name:>14}  {sparkline(sample(schedule, num_steps=total))}")

lrsched-0.1.0/pyproject.toml

@@ -0,0 +1,61 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "lrsched"
+version = "0.1.0"
+description = "Framework-agnostic learning-rate schedules as pure functions, in Python with zero dependencies."
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE" }
+authors = [{ name = "Amaar Chughtai" }]
+keywords = [
+  "learning-rate",
+  "scheduler",
+  "lr-schedule",
+  "cosine-annealing",
+  "warmup",
+  "one-cycle",
+  "sgdr",
+  "machine-learning",
+  "deep-learning",
+  "training",
+]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Science/Research",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Typing :: Typed",
+]
+dependencies = []
+
+[project.optional-dependencies]
+dev = ["pytest>=8", "ruff>=0.6", "mypy>=1.11", "hypothesis>=6"]
+
+[project.urls]
+Homepage = "https://github.com/amaar-mc/lrsched"
+Repository = "https://github.com/amaar-mc/lrsched"
+Issues = "https://github.com/amaar-mc/lrsched/issues"
+Changelog = "https://github.com/amaar-mc/lrsched/blob/main/CHANGELOG.md"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/lrsched"]
+
+[tool.ruff]
+line-length = 100
+src = ["src", "tests"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "SIM", "RUF"]
+
+[tool.mypy]
+strict = true
+files = ["src"]

lrsched-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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-0.1.0/src/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/tests/test_compose.py

@@ -0,0 +1,30 @@
+import pytest
+
+from lrsched import constant, cosine, sample, sequential, with_warmup
+
+
+def test_with_warmup() -> None:
+    s = with_warmup(cosine(base_lr=1.0, min_lr=0.0, total_steps=10), warmup_steps=5, start_lr=0.0)
+    assert s(0) == pytest.approx(0.0)
+    assert s(2) == pytest.approx(0.4)  # linear ramp 0 -> 1 over 5 steps
+    assert s(5) == pytest.approx(1.0)  # hands off to cosine at its step 0
+    assert s(10) == pytest.approx(cosine(base_lr=1.0, min_lr=0.0, total_steps=10)(5))
+
+
+def test_sequential() -> None:
+    s = sequential([(5, constant(lr=0.1)), (5, constant(lr=0.2))])
+    assert s(0) == pytest.approx(0.1)
+    assert s(4) == pytest.approx(0.1)
+    assert s(5) == pytest.approx(0.2)
+    assert s(9) == pytest.approx(0.2)
+    assert s(50) == pytest.approx(0.2)  # held at the last phase past the end
+
+
+def test_sequential_requires_a_phase() -> None:
+    with pytest.raises(ValueError):
+        sequential([])
+
+
+def test_sample() -> None:
+    assert sample(constant(lr=0.5), num_steps=3) == [0.5, 0.5, 0.5]
+    assert len(sample(cosine(base_lr=1.0, min_lr=0.0, total_steps=100), num_steps=100)) == 100

lrsched-0.1.0/tests/test_properties.py

@@ -0,0 +1,31 @@
+from itertools import pairwise
+
+from hypothesis import given
+from hypothesis import strategies as st
+
+from lrsched import constant, cosine, sample, with_warmup
+
+steps = st.integers(min_value=1, max_value=1000)
+rates = st.floats(min_value=0.0, max_value=10.0, allow_nan=False, allow_infinity=False)
+
+
+@given(rates, rates, steps, st.integers(min_value=0, max_value=2000))
+def test_cosine_stays_within_bounds(base: float, gap: float, total: int, step: int) -> None:
+    base_lr = base + gap
+    min_lr = base
+    value = cosine(base_lr=base_lr, min_lr=min_lr, total_steps=total)(step)
+    assert min_lr - 1e-9 <= value <= base_lr + 1e-9
+
+
+@given(steps, rates)
+def test_warmup_is_monotone_nondecreasing(warmup: int, target: float) -> None:
+    main = cosine(base_lr=target, min_lr=0.0, total_steps=100)
+    s = with_warmup(main, warmup_steps=warmup, start_lr=0.0)
+    values = [s(i) for i in range(warmup + 1)]
+    for a, b in pairwise(values):
+        assert b >= a - 1e-9
+
+
+@given(st.integers(min_value=1, max_value=500))
+def test_sample_length(num_steps: int) -> None:
+    assert len(sample(constant(lr=0.1), num_steps=num_steps)) == num_steps

lrsched-0.1.0/tests/test_schedules.py

@@ -0,0 +1,107 @@
+import pytest
+
+from lrsched import (
+    constant,
+    cosine,
+    cosine_restarts,
+    exponential,
+    inverse_sqrt,
+    linear,
+    multi_step,
+    one_cycle,
+    polynomial,
+    step_decay,
+)
+
+
+def test_constant() -> None:
+    s = constant(lr=0.1)
+    assert s(0) == 0.1
+    assert s(1000) == 0.1
+
+
+def test_exponential() -> None:
+    s = exponential(base_lr=1.0, gamma=0.9)
+    assert s(0) == pytest.approx(1.0)
+    assert s(1) == pytest.approx(0.9)
+    assert s(2) == pytest.approx(0.81)
+
+
+def test_step_decay() -> None:
+    s = step_decay(base_lr=1.0, gamma=0.5, step_size=10)
+    assert s(0) == pytest.approx(1.0)
+    assert s(9) == pytest.approx(1.0)
+    assert s(10) == pytest.approx(0.5)
+    assert s(20) == pytest.approx(0.25)
+
+
+def test_multi_step() -> None:
+    s = multi_step(base_lr=1.0, milestones=[20, 10], gamma=0.1)
+    assert s(9) == pytest.approx(1.0)
+    assert s(10) == pytest.approx(0.1)
+    assert s(19) == pytest.approx(0.1)
+    assert s(20) == pytest.approx(0.01)
+
+
+def test_linear() -> None:
+    s = linear(base_lr=1.0, end_lr=0.0, total_steps=10)
+    assert s(0) == pytest.approx(1.0)
+    assert s(5) == pytest.approx(0.5)
+    assert s(10) == pytest.approx(0.0)
+    assert s(50) == pytest.approx(0.0)  # held past the end
+
+
+def test_polynomial() -> None:
+    s = polynomial(base_lr=1.0, end_lr=0.0, total_steps=10, power=2.0)
+    assert s(0) == pytest.approx(1.0)
+    assert s(5) == pytest.approx(0.25)
+    assert s(10) == pytest.approx(0.0)
+
+
+def test_cosine() -> None:
+    s = cosine(base_lr=1.0, min_lr=0.0, total_steps=10)
+    assert s(0) == pytest.approx(1.0)
+    assert s(5) == pytest.approx(0.5)
+    assert s(10) == pytest.approx(0.0)
+    assert s(100) == pytest.approx(0.0)  # held at min past the end
+
+
+def test_cosine_restarts_periodic() -> None:
+    s = cosine_restarts(base_lr=1.0, min_lr=0.0, period=10, t_mult=1)
+    assert s(0) == pytest.approx(1.0)
+    assert s(5) == pytest.approx(0.5)
+    assert s(10) == pytest.approx(1.0)  # restart
+    assert s(15) == pytest.approx(0.5)
+
+
+def test_cosine_restarts_growing() -> None:
+    s = cosine_restarts(base_lr=1.0, min_lr=0.0, period=10, t_mult=2)
+    assert s(10) == pytest.approx(1.0)  # restart, second cycle has length 20
+    assert s(20) == pytest.approx(0.5)  # halfway through the length-20 cycle
+    assert s(30) == pytest.approx(1.0)  # restart, third cycle
+
+
+def test_inverse_sqrt() -> None:
+    s = inverse_sqrt(base_lr=1.0, warmup_steps=10)
+    assert s(0) == pytest.approx(0.0)
+    assert s(5) == pytest.approx(0.5)
+    assert s(10) == pytest.approx(1.0)  # peak at end of warmup
+    assert s(40) == pytest.approx(0.5)  # sqrt(10/40)
+
+
+def test_one_cycle() -> None:
+    s = one_cycle(max_lr=1.0, total_steps=100, pct_start=0.3, initial_div=25.0, final_div=10000.0)
+    assert s(0) == pytest.approx(0.04)  # max_lr / initial_div
+    assert s(30) == pytest.approx(1.0)  # peak at end of the ramp
+    assert s(100) == pytest.approx(0.0001)  # max_lr / final_div
+
+
+def test_validation() -> None:
+    with pytest.raises(ValueError):
+        constant(lr=float("inf"))
+    with pytest.raises(ValueError):
+        cosine(base_lr=1.0, min_lr=0.0, total_steps=0)
+    with pytest.raises(ValueError):
+        one_cycle(max_lr=1.0, total_steps=100, pct_start=1.5, initial_div=25.0, final_div=1e4)
+    with pytest.raises(ValueError):
+        constant(lr=0.1)(-1)