Skip to content

Commit b5da8b8

Browse files
feat(decoration)!: lazy wrapping — import torchlens has no side effects
BREAKING CHANGE: torch functions are no longer wrapped at import time. Wrapping happens lazily on first log_forward_pass() call and persists. Three changes: 1. Lazy decoration: removed decorate_all_once() / patch_detached_references() calls from __init__.py. _ensure_model_prepared() triggers wrapping on first use via wrap_torch(). 2. Public wrap/unwrap API: - torchlens.wrap_torch() — install wrappers (idempotent) - torchlens.unwrap_torch() — restore original torch callables - torchlens.wrapped() — context manager (wrap on enter, unwrap on exit) - log_forward_pass(unwrap_when_done=True) — one-shot convenience Old names (undecorate_all_globally, redecorate_all_globally) kept as internal aliases. 3. torch.identity fix: decorated identity function now stored on _state._decorated_identity instead of monkey-patching torch.identity (which doesn't exist in PyTorch type stubs). Eliminates 2 mypy errors. Tests updated: 75 pass including 12 new lifecycle tests.
1 parent fcb912c commit b5da8b8

File tree

7 files changed

+364
-116
lines changed

7 files changed

+364
-116
lines changed

0 commit comments

Comments
 (0)