Skip to content

Export SparseCalibrationWeights and add optional seed parameter#43

Merged
MaxGhenis merged 1 commit intomainfrom
fix/api-exports-and-seeds
Apr 17, 2026
Merged

Export SparseCalibrationWeights and add optional seed parameter#43
MaxGhenis merged 1 commit intomainfrom
fix/api-exports-and-seeds

Conversation

@MaxGhenis
Copy link
Copy Markdown
Contributor

Summary

  • Export SparseCalibrationWeights from the top-level l0 package — it was already documented in the README and in CRITICAL_TEMPERATURE_BUG.md as the temperature-correct path for calibration but was missing from __all__. SparseL0Linear was already exported.
  • Add optional seed: int | None kwarg to SparseCalibrationWeights and SparseL0Linear. When set, log_alpha jitter (and log_weight jitter inside SparseCalibrationWeights.fit) draws from a local torch.Generator; seed=None preserves the legacy global-RNG behaviour.
  • Update CLAUDE.md to list calibration.py / sparse.py as first-class modules, update test file layout, and replace the stale black -l 79 formatter hint with the current ruff format invocation.

Addresses bug-hunt findings

Test plan

  • uv run pytest tests -x -q passes (88 passed, 1 skipped; previously 83).
  • test_sparse_calibration_weights_exported pins "SparseCalibrationWeights" in l0.__all__.
  • test_seed_produces_deterministic_log_alpha and test_seed_none_uses_global_rng verify both RNG paths in test_calibration.py and test_sparse.py.

- `SparseCalibrationWeights` was documented in the README and in
  `CRITICAL_TEMPERATURE_BUG.md` as the temperature-correct path for
  calibration workflows, but it was not re-exported from `l0/__init__.py`.
  `SparseL0Linear` was already exported, so the asymmetry made
  `SparseCalibrationWeights` noticeably harder to discover. Adds it to the
  package's `__all__`.
- `SparseCalibrationWeights.__init__` and `SparseL0Linear.__init__` drew
  their `log_alpha` jitter from PyTorch's global RNG, so two constructions
  with the same inputs were only deterministic if the caller remembered to
  `torch.manual_seed(...)` beforehand. Adds an optional `seed: int | None`
  kwarg to both classes; when set, the jitter (and `log_weight` jitter
  inside `SparseCalibrationWeights.fit`) draws from a local
  `torch.Generator`. `seed=None` preserves legacy behaviour.
- Updates `CLAUDE.md` to list `calibration.py` and `sparse.py` as
  first-class modules, fixes the stale `black -l 79` formatter hint, and
  adds the matching test files.

Adds seed/export regression tests in `test_calibration.py` and
`test_sparse.py`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Copy link
Copy Markdown
Contributor Author

@MaxGhenis MaxGhenis left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

Verified:

  • l0/__init__.py: SparseCalibrationWeights added to __all__ alongside SparseL0Linear (previously only the latter was exported). Import works: from l0 import SparseCalibrationWeights succeeds; l0.SparseCalibrationWeights is SparseCalibrationWeights.
  • seed: int | None = None on both SparseCalibrationWeights.__init__ and SparseL0Linear.__init__, with matching docstring entries.
  • Default None preserves legacy global-RNG behaviour — test_seed_none_uses_global_rng pins this with torch.manual_seed(0) on both sides.
  • When seed is set, a local torch.Generator is created (on the module's device) and used consistently in both init jitter and (for SparseCalibrationWeights) the fit log_weight jitter. The randn_like-has-no-generator edge case is handled explicitly via torch.randn with shape + generator.
  • SparseL0Linear.__init__ also rewrites torch.normal(mu.item(), 0.01, size=...) to mu.item() + 0.01 * torch.randn(..., generator=...) so the seeded path has a single consistent RNG — nice tidy.
  • CLAUDE.md now lists calibration.py / sparse.py as first-class modules, updates the package tree, adds test_calibration.py / test_sparse.py, and replaces the stale black -l 79 formatter hint with ruff format. Subsumes finding #5's doc staleness.
  • 28 tests pass locally.

Will go green once #44 merges.

@MaxGhenis MaxGhenis merged commit 2a1e559 into main Apr 17, 2026
2 of 3 checks passed
@MaxGhenis MaxGhenis deleted the fix/api-exports-and-seeds branch April 17, 2026 16:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant