feat(decoration): add global undecorate override

Author: johnmarktaylor91

TODO.md

@@ -0,0 +1,6 @@
+# TorchLens TODO
+
+- Add a first-class "undecorate / suspend decoration" escape hatch so users can
+  restore a clean PyTorch environment when tracing is not needed.
+- Explore decorating on the first `log_forward_pass(...)` instead of at import
+  time, while still leaving torch decorated afterward once the user opts in.

tests/test_decoration.py

@@ -22,6 +22,8 @@
     decorate_all_once,
     patch_detached_references,
     patch_model_instance,
+    redecorate_all_globally,
+    undecorate_all_globally,
 )
 
 
@@ -179,6 +181,33 @@ def test_requires_grad_restored_after_exception(self):
             assert param.requires_grad == orig_grads[name], f"{name} requires_grad changed"
 
 
+class TestGlobalUndecorate:
+    def test_undecorate_restores_original_torch_function(self):
+        """Users can globally strip TorchLens wrappers from torch callables."""
+        assert getattr(torch.cos, "tl_is_decorated_function", False)
+        undecorate_all_globally()
+        try:
+            assert not getattr(torch.cos, "tl_is_decorated_function", False)
+            x = torch.randn(4)
+            y = torch.cos(x)
+            assert y.shape == x.shape
+            assert not hasattr(y, "tl_tensor_label_raw")
+        finally:
+            redecorate_all_globally()
+
+    def test_redecorate_restores_logging_after_global_undecorate(self):
+        """Global undecoration is reversible for later TorchLens use."""
+        undecorate_all_globally()
+        try:
+            assert not getattr(torch.nn.functional.relu, "tl_is_decorated_function", False)
+        finally:
+            redecorate_all_globally()
+        assert getattr(torch.nn.functional.relu, "tl_is_decorated_function", False)
+        result = log_forward_pass(SimpleModel(), torch.randn(5))
+        relu_layers = [label for label in result.layer_labels if "relu" in label.lower()]
+        assert relu_layers
+
+
 # =========================================================================
 # 2. Torch Functions Normal When Toggle Off
 # =========================================================================

torchlens/CLAUDE.md

@@ -66,6 +66,8 @@ class definition and the corresponding FIELD_ORDER in constants.py.
 1. `_state.py` must never import other torchlens modules
 2. RNG state capture/restore must happen BEFORE `active_logging()` context
 3. `pause_logging()` must wrap any internal torch ops during logging (safe_copy, activation_postfunc, get_tensor_memory_amount)
-4. Decorated wrappers are permanent — never undecorated
+4. Decorated wrappers are permanent by default, but advanced users may call
+   `undecorate_all_globally()` / `redecorate_all_globally()` as an explicit
+   override when they need a clean PyTorch environment.
 5. Field-order constants and class definitions must stay in sync
 6. Step 6 module suffix mutation makes `_rebuild_pass_assignments` (Step 8) NECESSARY — not just defensive

torchlens/__init__.py

@@ -30,6 +30,7 @@
 from .data_classes.layer_pass_log import LayerPassLog, TensorLog
 from .data_classes import FuncCallLocation, ModuleAccessor, ModuleLog, ModulePassLog, ParamLog
 from .visualization import build_render_audit, model_log_to_dagua_graph, render_model_log_with_dagua
+from .decoration import redecorate_all_globally, undecorate_all_globally
 
 # ---- Import-time decoration (side effects) --------------------------------
 #

torchlens/decoration/__init__.py

@@ -1 +1,17 @@
 """Torch function wrapping and model preparation for logging."""
+
+from .torch_funcs import (
+    decorate_all_once,
+    patch_detached_references,
+    patch_model_instance,
+    redecorate_all_globally,
+    undecorate_all_globally,
+)
+
+__all__ = [
+    "decorate_all_once",
+    "patch_detached_references",
+    "patch_model_instance",
+    "redecorate_all_globally",
+    "undecorate_all_globally",
+]

torchlens/decoration/torch_funcs.py

@@ -1,4 +1,4 @@
-"""Permanent torch function wrapping: decorates all torch ops at import time with toggle-gated wrappers.
+"""Permanent-by-default torch function wrapping with optional global override.
 
 This module implements the core interception mechanism for TorchLens. At ``import torchlens``
 time, every function listed in ``ORIG_TORCH_FUNCS`` is replaced with a thin wrapper that
@@ -9,8 +9,11 @@
 
 Key design decisions:
 
-1. **Permanent decoration** avoids the fragility of repeatedly patching/unpatching torch
-   internals. The toggle makes this safe for production use.
+1. **Permanent decoration by default** avoids the fragility of repeatedly
+   patching/unpatching torch internals. The toggle makes this safe for production
+   use.  Advanced users can still call ``undecorate_all_globally()`` to restore
+   the original torch callables and ``redecorate_all_globally()`` to re-enable
+   TorchLens interception.
 
 2. **Shared originals reuse wrappers**: If ``torch.cos`` and ``torch._VF.cos`` point to the
    same C builtin, only one wrapper is created and both namespaces point to it. This keeps
@@ -584,6 +587,131 @@ def decorate_all_once():
     torch.identity = new_identity
 
 
+def _replace_detached_references(mapping: Dict[int, Callable]) -> None:
+    """Crawl ``sys.modules`` and replace callable references using ``mapping``.
+
+    ``mapping`` may be either original->decorated or decorated->original. This
+    keeps the sys.modules crawl logic symmetric so ``undecorate`` can reverse
+    the same module/class/default-arg patching done during normal decoration.
+    """
+    if not mapping:
+        return
+
+    for mod in list(sys.modules.values()):
+        if mod is None:
+            continue
+        if hasattr(mod, "__name__") and getattr(mod, "__name__", "").startswith("torchlens"):
+            continue
+
+        try:
+            mod_dict = vars(mod)
+        except TypeError:
+            continue
+
+        for attr_name, attr_val in list(mod_dict.items()):
+            if id(attr_val) in mapping:
+                try:
+                    mod_dict[attr_name] = mapping[id(attr_val)]
+                except (TypeError, KeyError):
+                    pass
+                continue
+
+            try:
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore")
+                    is_type = isinstance(attr_val, type)
+            except Exception:
+                is_type = False
+            if is_type:
+                try:
+                    cls_dict = vars(attr_val)
+                except TypeError:
+                    continue
+                for cls_attr_name, cls_attr_val in list(cls_dict.items()):
+                    if id(cls_attr_val) in mapping:
+                        try:
+                            setattr(attr_val, cls_attr_name, mapping[id(cls_attr_val)])
+                        except (AttributeError, TypeError):
+                            pass
+
+            try:
+                is_callable = callable(attr_val) and not is_type
+            except Exception:
+                is_callable = False
+            if is_callable:
+                _patch_function_defaults(attr_val, mapping)
+
+
+def undecorate_all_globally() -> None:
+    """Restore original torch callables globally.
+
+    This is an explicit override for advanced users who want a clean PyTorch
+    environment after importing TorchLens. It restores torch namespace
+    attributes and reverses the detached-reference crawl so module globals,
+    class attributes, and function defaults point back at original callables.
+
+    TorchLens logging will not function until ``redecorate_all_globally()`` is
+    called again.
+    """
+    _state._logging_enabled = False
+    _state._active_model_log = None
+
+    if not _state._decorated_to_orig:
+        return
+
+    for namespace_name, func_name in ORIG_TORCH_FUNCS:
+        namespace_key = namespace_name.replace("torch.", "")
+        local_func_namespace = nested_getattr(torch, namespace_key)
+        if not hasattr(local_func_namespace, func_name):
+            continue
+        current = getattr(local_func_namespace, func_name)
+        orig = _state._decorated_to_orig.get(id(current))
+        if orig is None:
+            continue
+        try:
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                setattr(local_func_namespace, func_name, orig)
+        except (AttributeError, TypeError):
+            pass
+
+    _replace_detached_references(_state._decorated_to_orig)
+    torch.identity = identity
+
+
+def redecorate_all_globally() -> None:
+    """Reinstall TorchLens wrappers after a prior global undecoration."""
+    if not _state._orig_to_decorated:
+        decorate_all_once()
+        patch_detached_references()
+        return
+
+    for namespace_name, func_name in ORIG_TORCH_FUNCS:
+        namespace_key = namespace_name.replace("torch.", "")
+        local_func_namespace = nested_getattr(torch, namespace_key)
+        if not hasattr(local_func_namespace, func_name):
+            continue
+        current = getattr(local_func_namespace, func_name)
+        decorated = None
+        if id(current) in _state._orig_to_decorated:
+            decorated = _state._orig_to_decorated[id(current)]
+        elif id(current) in _state._decorated_to_orig:
+            decorated = current
+        if decorated is None:
+            continue
+        try:
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                setattr(local_func_namespace, func_name, decorated)
+        except (AttributeError, TypeError):
+            pass
+
+    _replace_detached_references(_state._orig_to_decorated)
+    new_identity = torch_func_decorator(identity, "identity")
+    new_identity.tl_is_decorated_function = True
+    torch.identity = new_identity
+
+
 # ---------------------------------------------------------------------------
 # sys.modules deep crawl
 # ---------------------------------------------------------------------------