Commit b5da8b8
committed
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- tests
- torchlens
- decoration
7 files changed
+364
-116
lines changed
0 commit comments