Decorators

Overview

Decorators are core to how PyTorch works:

• @torch.jit.script = Compile function to TorchScript
• @property = Lazy evaluation of model attributes
• @abstractmethod = Define interfaces
• Custom decorators = Add functionality to layers, profiling, debugging

Decorators are "functions that take functions and return functions."

---

Decorator Basics

Simple Decorator

def my_decorator(func):
 """Wraps a function."""
 def wrapper(*args, **kwargs):
 print(f"Before calling {func.__name__}")
 result = func(*args, **kwargs)
 print(f"After calling {func.__name__}")
 return result
 return wrapper

@my_decorator # Equivalent to: add = my_decorator(add)
def add(a, b):
 return a + b

add(1, 2)
# Output:
# Before calling add
# After calling add

How Decorators Work

# These are equivalent:

# Using @ syntax
@my_decorator
def add(a, b):
 return a + b

# Explicit call
def add(a, b):
 return a + b
add = my_decorator(add)

---

PyTorch Decorators

@property: Lazy Evaluation

class Tensor:
 def __init__(self, data):
 self._data = data
 self._shape = None # Computed lazily

 @property
 def shape(self):
 """Compute shape on demand."""
 if self._shape is None:
 self._shape = tuple(len(d) for d in self._data)
 return self._shape

t = Tensor([[1, 2], [3, 4]])
print(t.shape) # (2, 2) - computed when accessed

@staticmethod and @classmethod

class Tensor:
 count = 0 # Class variable

 def __init__(self, data):
 self.data = data
 Tensor.count += 1

 @staticmethod
 def create_zeros(shape):
 """Static method (no self)."""
 return Tensor([0] * shape[0])

 @classmethod
 def from_numpy(cls, array):
 """Class method (has cls)."""
 return cls(array.tolist())

# Usage
t1 = Tensor.create_zeros((3,))
t2 = Tensor.from_numpy(np.array([1, 2, 3]))
print(Tensor.count) # 2

@abstractmethod

from abc import ABC, abstractmethod

class Layer(ABC):
 @abstractmethod
 def forward(self, x):
 """Must be implemented by subclasses."""
 pass

class Linear(Layer):
 def forward(self, x):
 return x @ self.weight

# Linear()(x) works - forward implemented
# Layer()(x) fails - abstract method not implemented

---

Custom Decorators for ML

Timing Decorator

import functools
import time

def timing_decorator(func):
 """Measure function execution time."""
 @functools.wraps(func) # Preserve function metadata
 def wrapper(*args, **kwargs):
 start = time.time()
 result = func(*args, **kwargs)
 elapsed = time.time() - start
 print(f"{func.__name__} took {elapsed:.3f}s")
 return result
 return wrapper

@timing_decorator
def train_epoch(model, data):
 # Training code...
 return loss

train_epoch(model, data)
# Output

Validation Decorator

def validate_inputs(func):
 """Validate tensor inputs."""
 @functools.wraps(func)
 def wrapper(self, x):
 if not isinstance(x, torch.Tensor):
 raise TypeError(f"Expected Tensor, got {type(x)}")
 if x.dim() == 0:
 raise ValueError("Tensor cannot be 0-dimensional")
 return func(self, x)
 return wrapper

class Model(torch.nn.Module):
 @validate_inputs
 def forward(self, x):
 return self.fc(x)

model = Model(torch.nn.Linear(10, 5))
model(torch.randn(3, 10)) # OK
model(42) # TypeError: Expected Tensor
model(torch.tensor(1)) # ValueError: Tensor cannot be 0-dimensional

Gradient Tracking Decorator

def with_grad_tracking(func):
 """Track gradients for debugging."""
 @functools.wraps(func)
 def wrapper(self, *args, **kwargs):
 results = []
 for i, arg in enumerate(args):
 if isinstance(arg, torch.Tensor) and arg.requires_grad:
 print(f"Input {i} requires grad")

 output = func(self, *args, **kwargs)

 if isinstance(output, torch.Tensor) and output.requires_grad:
 print(f"Output requires grad")

 return output
 return wrapper

class Model(torch.nn.Module):
 @with_grad_tracking
 def forward(self, x):
 return self.fc(x)

---

Decorators with Arguments

Fixed Arguments

def repeat_n_times(n):
 """Repeat function call n times."""
 def decorator(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 results = []
 for _ in range(n):
 result = func(*args, **kwargs)
 results.append(result)
 return results
 return wrapper
 return decorator

@repeat_n_times(3)
def create_tensor():
 return torch.randn(3)

tensors = create_tensor() # Creates 3 tensors
print(len(tensors)) # 3

Variable Arguments

def retry(max_attempts=3, delay=1):
 """Retry function on failure."""
 def decorator(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 import time
 for attempt in range(max_attempts):
 try:
 return func(*args, **kwargs)
 except Exception as e:
 if attempt == max_attempts - 1:
 raise
 print(f"Attempt {attempt+1} failed, retrying in {delay}s...")
 time.sleep(delay)
 return wrapper
 return decorator

@retry(max_attempts=3, delay=2)
def load_model(path):
 return torch.load(path)

model = load_model("model.pt")

---

Chaining Decorators

Multiple decorators applied bottom-to-top:

@timing_decorator
@validate_inputs
def forward(self, x):
 return self.fc(x)

# Equivalent to:
def forward(self, x):
 return self.fc(x)
forward = validate_inputs(forward)
forward = timing_decorator(forward)

# When called:
# 1. timing_decorator wrapper called
# 2. Inside, calls validate_inputs wrapper
# 3. Inside, calls original forward
# 4. Returns to validate_inputs wrapper, then timing_decorator wrapper

Practical: Combined Decorators

def debug(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 print(f"Calling {func.__name__}")
 return func(*args, **kwargs)
 return wrapper

def validate(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 if len(args) == 0:
 raise ValueError("Need arguments")
 return func(*args, **kwargs)
 return wrapper

@debug
@validate
def forward(x):
 return x * 2

forward(5)
# Output:
# Calling forward
# 10

forward() # Raises ValueError before debug prints

---

Class Decorators

Decorator on Class

def add_repr(cls):
 """Add repr to any class."""
 def __repr__(self):
 attrs = ", ".join(f"{k}={v}" for k, v in self.__dict__.items())
 return f"{cls.__name__}({attrs})"

 cls.__repr__ = __repr__
 return cls

@add_repr
class Point:
 def __init__(self, x, y):
 self.x = x
 self.y = y

p = Point(3, 4)
print(p) # Point(x=3, y=4)

Registering Classes

# Registry pattern (used in PyTorch)
MODEL_REGISTRY = {}

def register_model(name):
 def decorator(cls):
 MODEL_REGISTRY[name] = cls
 return cls
 return decorator

@register_model("resnet50")
class ResNet50(torch.nn.Module):
 pass

@register_model("vgg16")
class VGG16(torch.nn.Module):
 pass

# Later
model_class = MODEL_REGISTRY["resnet50"]
model = model_class()

---

Advanced: Functools.wraps

Preserves function metadata:

import functools

def bad_decorator(func):
 def wrapper(*args, **kwargs):
 return func(*args, **kwargs)
 return wrapper

def good_decorator(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 return func(*args, **kwargs)
 return wrapper

def add(x, y):
 """Add two numbers."""
 return x + y

bad = bad_decorator(add)
good = good_decorator(add)

print(bad.__name__) # 'wrapper' (lost original)
print(good.__name__) # 'add' (preserved)

print(bad.__doc__) # None
print(good.__doc__) # 'Add two numbers.'

---

Decorator Stacking Performance

import timeit

def timer_decorator(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 return func(*args, **kwargs)
 return wrapper

# Single decorator
@timer_decorator
def func1(x):
 return x * 2

# Multiple decorators
@timer_decorator
@timer_decorator
@timer_decorator
def func3(x):
 return x * 2

time1 = timeit.timeit(lambda: func1(42), number=1000000)
time3 = timeit.timeit(lambda: func3(42), number=1000000)

print(f"Single: {time1:.3f}s")
print(f"Triple: {time3:.3f}s")
# Overhead is minimal (<5%)

---

PyTorch-Specific Decorators

TorchScript Compilation

@torch.jit.script
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 return x + y

# Compiled to TorchScript (faster, deployable)
x = torch.tensor([1.0, 2.0])
y = torch.tensor([3.0, 4.0])
print(add(x, y))

Custom Gradient

class CustomFunction(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x):
 ctx.save_for_backward(x)
 return x ** 2

 @staticmethod
 def backward(ctx, grad_output):
 x, = ctx.saved_tensors
 return 2 * x * grad_output

# Use in model
x = torch.tensor([2.0], requires_grad=True)
y = CustomFunction.apply(x)
y.backward()
print(x.grad) # tensor([4.])

---

Common Pitfalls

Pitfall 1: Forgetting functools.wraps

# BAD
def bad_decorator(func):
 def wrapper(*args, **kwargs):
 return func(*args, **kwargs)
 return wrapper

# GOOD
def good_decorator(func):
 @functools.wraps(func)
 def wrapper(*args, **kwargs):
 return func(*args, **kwargs)
 return wrapper

Pitfall 2: Decorators Changing Function Signature

# This hides the actual signature
def decorator(func):
 def wrapper(*args, **kwargs):
 return func(*args, **kwargs)
 return wrapper

@decorator
def add(a, b):
 return a + b

# IDE can't show add(a, b) signature
# Instead shows wrapper(*args, **kwargs)

---

Summary

• Decorators = Functions that modify functions
• @property = Make methods look like attributes
• @staticmethod = Function bound to class, not instance
• @classmethod = Method receives class as first argument
• Custom decorators = Add functionality (timing, validation, caching)
• Chaining = Multiple decorators compose
• functools.wraps = Always use to preserve metadata

-

Related Topics

• [01 Classes & Inheritance](/05-py3/02-object-oriented-patterns/(01-classes-inheritance/) - Decorators work with classes too
• [04 Descriptors & Properties](/05-py3/02-object-oriented-patterns/(04-descriptors-properties/) - Advanced property mechanisms
• 00 Readme - Higher-order functions