Type System & Annotations

Overview

Python's dynamic type system combined with optional static type hints creates the perfect balance for ML:

• Rapid prototyping without type declarations
• IDE autocomplete with type annotations
• Runtime validation with type checking

This is fundamental to why PyTorch and JAX can be both:

• Easy to use (dynamic typing)
• Maintainable at scale (type hints)

-

Dynamic Typing: The Core

Python uses runtime type checking, not compile-time:

# Same variable can be different types
x = 5 # int
x = "tensor" # str
x = [1, 2, 3] # list

Why this matters for ML:

• Enables flexible tensor operations (NumPy array protocol)
• Allows framework magic (PyTorch's autograd hook)
• No compilation step = faster iteration

Type Introspection at Runtime

import sys
from torch import Tensor

# Check type at runtime
x = Tensor([1, 2, 3])
print(type(x)) # <class 'torch.Tensor'>
print(isinstance(x, Tensor)) # True

# Get attributes and methods dynamically
print(dir(x)) # All available methods
print(hasattr(x, 'shape')) # True
print(callable(x.backward)) # True

Framework Pattern: PyTorch uses isinstance checks to dispatch between CPU/CUDA/MPS tensors at runtime.

---

Type Hints: Static Analysis

Added in Python 3.5, type hints are optional annotations that don't affect runtime:

from typing import List, Dict, Tuple, Optional
import torch

# Function type hints
def forward(
 input: torch.Tensor,
 weights: torch.Tensor,
 bias: Optional[torch.Tensor] = None
) -> torch.Tensor:
 """Compute linear transformation."""
 output = input @ weights.T
 if bias is not None:
 output += bias
 return output

# Variable annotations
model: torch.nn.Module = torch.nn.Linear(10, 20)
embeddings: Dict[str, torch.Tensor] = {}

Why PyTorch recommends type hints:

• IDE autocomplete for .shape, .dtype, methods
• Static analysis catches errors before runtime
• Documentation for users
• Enables type checkers (mypy, pyright)

Benefits for ML Frameworks

# Without type hints
def create_tensor(shape, dtype):
 return torch.zeros(shape, dtype=dtype)

# With type hints - IDE knows what to suggest
def create_tensor(shape: Tuple[int,...], 
 dtype: torch.dtype = torch.float32) -> torch.Tensor:
 return torch.zeros(shape, dtype=dtype)

---

Built-in Types vs Framework Types

Type Hierarchy in Python

# Abstract Base Classes
from collections.abc import Sequence, Mapping, Callable

# NumPy arrays follow sequence protocol
import numpy as np
arr = np.array([1, 2, 3])
print(isinstance(arr, Sequence)) # False - NumPy doesn't inherit from ABC
print(hasattr(arr, '__getitem__')) # True - but duck typing

# Type checking protocols (Python 3.8+)
from typing import Protocol

class HasShape(Protocol):
 """Anything with a shape attribute."""
 shape: Tuple[int,...]

def print_shape(obj: HasShape) -> None:
 print(obj.shape)

# Both work without inheritance
print_shape(torch.zeros(3, 4)) # torch.Tensor
print_shape(np.zeros((3, 4))) # np.ndarray

Framework Design: JAX and PyTorch use protocols to accept "array-like" inputs without explicit type hierarchy.

-

Type Variables & Generics

Crucial for building reusable framework code:

from typing import TypeVar, Generic, List

# Generic type variable
T = TypeVar('T') # Can be any type

# Generic class (like PyTorch's Module)
class Container(Generic[T]):
 def __init__(self, items: List[T]):
 self.items = items

 def get(self, index: int) -> T:
 return self.items[index]

# Usage
int_container: Container[int] = Container([1, 2, 3])
tensor_container: Container[torch.Tensor] = Container([
 torch.zeros(3),
 torch.ones(3)
])

Constrained TypeVars

from typing import TypeVar, Union

# Only allow numeric types
Numeric = TypeVar('Numeric', int, float, complex)

def add(a: Numeric, b: Numeric) -> Numeric:
 return a + b

add(1, 2) # OK
add(1.5, 2.5) # OK
add("a", "b") # Type error in mypy

PyTorch Pattern: Linear layer uses generics for input/output tensor types.

---

Type Narrowing & Guards

Refining types based on runtime checks:

from typing import Union
import torch

def backward(tensor: Union[torch.Tensor, list]) -> None:
 # Type narrowing with isinstance
 if isinstance(tensor, torch.Tensor):
 tensor.backward() # Type checker knows it's Tensor
 else:
 # Type checker knows it's list here
 for t in tensor:
 t.backward()

# With TypeGuard (Python 3.10+)
from typing import TypeGuard

def is_tensor(x: object) -> TypeGuard[torch.Tensor]:
 return isinstance(x, torch.Tensor)

def process(x: Union[torch.Tensor, list]) -> None:
 if is_tensor(x):
 # Type checker narrows x to torch.Tensor
 print(x.shape)

---

Runtime Type Checking

Type hints are not enforced at runtime by default:

# No runtime error even though type is wrong
def forward(x: int) -> int:
 return x * 2

forward("hello") # Returns "hellohello" - no error!

Solution: Use runtime validators

from typing import get_type_hints

def validate_types(func):
 """Decorator that checks type hints at runtime."""
 hints = get_type_hints(func)

 def wrapper(*args, **kwargs):
 # Check arguments
 for i, (arg, (name, expected_type)) in enumerate(
 zip(args, hints.items())
):
 if not isinstance(arg, expected_type):
 raise TypeError(
 f"Argument {name} must be {expected_type}, "
 f"got {type(arg)}"
)
 return func(*args, **kwargs)

 return wrapper

@validate_types
def add(a: int, b: int) -> int:
 return a + b

add(1, 2) # OK
add(1, "2") # RuntimeError

Framework Pattern: PyTorch uses this for checking tensor shapes and dtypes in operations.

---

Type Checking with mypy

Using static type checker to find errors:

# Install mypy
pip install mypy

# Check types
mypy my_model.py

# mypy catches these errors
import torch

x: torch.Tensor = torch.zeros(3)
y: int = x # Error: Incompatible types

def forward(x: torch.Tensor) -> torch.Tensor:
 return x.shape # Error: Incompatible return type

# Correct version
def forward(x: torch.Tensor) -> torch.Size:
 return x.shape

ML Framework Setup:

# mypy.ini for PyTorch projects
[mypy]
python_version = 3.10
plugins = 
 numpy.typing.plugin
ignore_missing_imports = True # For untyped dependencies

[mypy-torch.*]
ignore_errors = True # Some torch internals untyped

-

Practical Patterns in PyTorch/JAX

Pattern 1: Device-Agnostic Tensors

from typing import Union
import torch

def forward(
 x: torch.Tensor,
 device: Union[str, torch.device] = "cpu"
) -> torch.Tensor:
 """Move tensor to device and process."""
 x = x.to(device)
 return x * 2

Pattern 2: Optional Tensors

from typing import Optional
import torch

class AttentionHead(torch.nn.Module):
 def forward(
 self,
 query: torch.Tensor,
 key: torch.Tensor,
 value: torch.Tensor,
 mask: Optional[torch.Tensor] = None
) -> torch.Tensor:
 """Self-attention with optional mask."""
 scores = query @ key.T
 if mask is not None:
 scores = scores.masked_fill(mask, float('-inf'))
 return torch.softmax(scores, dim=-1) @ value

Pattern 3: Union Types for Flexibility

from typing import Union, List
import torch

def stack_tensors(inputs: Union[torch.Tensor, List[torch.Tensor]]) -> torch.Tensor:
 """Accept single tensor or list of tensors."""
 if isinstance(inputs, list):
 return torch.stack(inputs)
 return inputs

# Both work
stack_tensors(torch.zeros(3))
stack_tensors([torch.zeros(3), torch.ones(3)])

---

Advanced: Protocol Types

Structural typing (duck typing with types):

from typing import Protocol, runtime_checkable
import torch
import numpy as np

@runtime_checkable
class TensorLike(Protocol):
 """Anything with shape, dtype, and array interface."""
 shape: tuple
 dtype: object

 def __array__(self) -> np.ndarray:
...

def print_info(tensor: TensorLike) -> None:
 print(f"Shape: {tensor.shape}, dtype: {tensor.dtype}")

# All work without explicit inheritance
print_info(torch.zeros(3, 4))
print_info(np.zeros((3, 4)))
print_info([1, 2, 3]) # Lists have shape? No - runtime error

---

Trade-offs: Dynamic vs Static Typing

Aspect	Dynamic	Static (with hints)
Development Speed	Fast prototyping	Slightly slower
IDE Support	Basic	Excellent
Error Detection	Runtime only	Before runtime
Refactoring Safety	Error-prone	Safe
Framework Magic	Enabled	Harder to implement
Learning Curve	Easy	Moderate

Python's Solution: Use dynamic typing for rapid prototyping, add type hints as code matures.

---

Summary

• Dynamic typing = Core Python strength for ML flexibility
• Type hints = Documentation + IDE support + type checking
• Runtime introspection = Framework magic (PyTorch hooks, JAX transformations)
• Protocols = Accept "anything array-like" without inheritance
• TypeVars = Build generic, reusable framework code

This combination is why Python dominates ML: maximum flexibility with optional type safety.

-

Related Topics

• 02 Data Structures Essentials - Built-in types optimized for ML
• [01 Getattr Setattr & Proxies](/05-py3/06-dynamic-features/(01-getattr-setattr-proxies/) - Runtime type hooking
• [02 Generics & Protocols](/05-py3/07-advanced-typing/(02-generics-protocols/) - Type system extensibility
• [02 Introspection & Reflection](/05-py3/06-dynamic-features/(02-introspection-reflection/) - type() and isinstance internals