Add annotated() decorator form of annotate()

`annotated("label")` wraps a whole function so every node it creates is
tagged with the label, equivalent to wrapping the body in `with
annotate(label):` but without re-indenting.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: physicsrob

torchwright/graph/__init__.py

@@ -1,4 +1,4 @@
-from .node import Node, annotate
+from .node import Node, annotate, annotated
 from .value_type import NodeValueType, Range
 from .attn import Attn
 from .embedding import Embedding

torchwright/graph/node.py

@@ -1,3 +1,4 @@
+import functools
 import os
 from contextlib import contextmanager
 from contextvars import ContextVar
@@ -58,6 +59,31 @@ def annotate(label: str):
         _current_annotation.reset(token)
 
 
+def annotated(label: str):
+    """Decorator form of :func:`annotate`.
+
+    Tags every node created during the wrapped function's execution with
+    ``label`` (hierarchical, same nesting rules as :func:`annotate`)::
+
+        @annotated("wall projection")
+        def project_segs(...):
+            ...
+
+    Equivalent to wrapping the whole body in ``with annotate(label):`` but
+    without re-indenting it.
+    """
+
+    def deco(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            with annotate(label):
+                return fn(*args, **kwargs)
+
+        return wrapper
+
+    return deco
+
+
 class Node:
     """Base class for all computation graph nodes.