Skip to content

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.