# Type Hints For Ml Code > Apply appropriate type hints for ML/PyTorch code. Use when adding type annotations to ML code or addressing mypy errors. Use when this capability is needed. - Skill: `tomevault-io/type-hints-for-ml-code` (Agent Skill, multi-file: 2 files) - Install (CLI): `npx skillmds@latest add tomevault-io/type-hints-for-ml-code` - Raw SKILL.md: https://api.skillmd.com/api/skills/tomevault-io/type-hints-for-ml-code/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - Author: tomevault-io (https://skillmd.com/u/tomevault-io) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/tomevault-io/type-hints-for-ml-code --- # Type Hints for ML Code - Constitutional AI ## When to Apply Automatically activate when: - Adding type hints to ML code - Addressing mypy type errors - Working with PyTorch tensors, models, or optimizers - Dealing with HuggingFace transformers types ## Project Type Checking Status **Current mypy status:** - **41 errors remaining** (documented in MYPY_ANALYSIS_REPORT.md) - **Status**: Accepted as reasonable for ML research code - **CI behavior**: `continue-on-error: true` (doesn't fail builds) **Key insight**: Perfect type coverage is not the goal for ML code. Prioritize correctness and readability over type perfection. ## When to Use Type Hints ### ✅ Do Use Type Hints For 1. **Public API functions** ```python def load_model(model_name: str) -> tuple[AutoModelForCausalLM, AutoTokenizer]: """Load model and tokenizer.""" pass ``` 2. **Configuration dataclasses** ```python @dataclass class TrainingConfig: learning_rate: float batch_size: int num_epochs: int ``` 3. **Clear input/output types** ```python def evaluate_text(text: str, framework: ConstitutionalFramework) -> dict[str, Any]: """Evaluate text against principles.""" pass ``` 4. **Helper functions with simple types** ```python def format_prompt(prompt: str, examples: list[str]) -> str: """Format prompt with examples.""" pass ``` ### ❌ Don't Force Type Hints For 1. **Complex tensor operations** (mypy struggles with tensor shapes) 2. **Dynamic PyTorch internals** (intentionally uses `Any`) 3. **NumPy operations** (overload resolution issues) 4. **Training loop internals** (too complex, low value) ## PyTorch Type Patterns ### Tensor Types ```python import torch from torch import Tensor # ✅ Basic tensor type def forward(inputs: Tensor) -> Tensor: return inputs * 2 # ✅ Optional tensor (common in ML) def process_batch( inputs: Tensor, labels: Tensor | None = None ) -> Tensor: pass # ⚠️ Avoid overly specific tensor types (mypy can't verify shapes) # This is too specific and won't type-check well: def bad_example(inputs: Tensor[int, 32, 768]) -> Tensor[int, 32, 10]: pass ``` ### Model and Optimizer Types ```python from torch.nn import Module from torch.optim import Optimizer from transformers import PreTrainedModel, PreTrainedTokenizer # ✅ Use base classes for flexibility def train_model( model: Module, # or PreTrainedModel for HuggingFace optimizer: Optimizer, dataloader: DataLoader ) -> None: pass # ✅ HuggingFace specific types def generate_text( prompt: str, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, max_length: int = 100 ) -> str: pass ``` ### Device Types ```python from torch import device as Device # ✅ Device type def move_to_device(tensor: Tensor, device: Device | str) -> Tensor: return tensor.to(device) # Common pattern: device: str | Device = torch.device("cuda" if torch.cuda.is_available() else "cpu") ``` ## Handling "Any" Type ### When `Any` is Acceptable ```python from typing import Any # ✅ For complex nested structures def process_model_output(output: Any) -> dict[str, Any]: """Process model output (structure varies by model).""" pass # ✅ For highly dynamic operations def advanced_tensor_op(tensors: list[Tensor]) -> Any: """Complex operation with unpredictable output type.""" pass # ✅ For configuration dictionaries config: dict[str, Any] = { 'learning_rate': 0.001, 'model_name': 'gpt2', 'device': 'cuda', } ``` ### Prefer Specific Types When Possible ```python # ❌ Too vague def process_data(data: Any) -> Any: pass # ✅ More specific def process_data(data: list[dict[str, float]]) -> dict[str, Tensor]: pass ``` ## Optional and Union Types ```python from typing import Optional # or use | None (Python 3.10+) # Modern syntax (Python 3.10+) def load_checkpoint(path: str | None = None) -> dict[str, Tensor] | None: if path is None: return None return torch.load(path) # Equivalent older syntax from typing import Optional, Union def load_checkpoint(path: Optional[str] = None) -> Optional[dict[str, Tensor]]: pass ``` ## Common ML Type Patterns ### Return Multiple Values ```python # ✅ Tuple with type hints def load_model(name: str) -> tuple[PreTrainedModel, PreTrainedTokenizer]: model = AutoModelForCausalLM.from_pretrained(name) tokenizer = AutoTokenizer.from_pretrained(name) return model, tokenizer # Usage model, tokenizer = load_model("gpt2") ``` ### Dataclass for Complex Returns ```python from dataclasses import dataclass @dataclass class EvaluationResult: any_flagged: bool flagged_principles: list[str] weighted_score: float details: dict[str, Any] def evaluate_text(text: str) -> EvaluationResult: """Evaluate text and return structured result.""" pass ``` ### Generator Types ```python from collections.abc import Iterator def batch_generator(data: list[Tensor], batch_size: int) -> Iterator[Tensor]: """Generate batches from data.""" for i in range(0, len(data), batch_size): yield data[i:i + batch_size] ``` ## Known Mypy Challenges in Project ### Category 1: NumPy Overload Resolution (14 errors) ```python # Mypy struggles with NumPy's 6+ overload variants import numpy as np # This may show mypy error, but it's correct at runtime scores = np.array([1.0, 2.0, 3.0]) # mypy: "cannot infer type" ``` **Resolution**: Use `# type: ignore[misc]` or accept the error (CI allows it) ### Category 2: Optional Attribute Access (7 errors) ```python # Mypy warns about potential None access def process_model(model: PreTrainedModel | None): if model is not None: output = model.generate(...) # mypy may still warn # Resolution: Assert non-None or use type: ignore ``` ### Category 3: Tensor Type Inference (8 errors) ```python # PyTorch intentionally uses dynamic typing loss = criterion(outputs, targets) # mypy: "Cannot determine type" ``` **Resolution**: Accept as limitation - PyTorch is dynamically typed by design ## Type Hints Best Practices for This Project ### 1. Prioritize Public APIs ```python # ✅ Type hints for exported functions def setup_default_framework() -> ConstitutionalFramework: """Public API - should have types.""" pass # ⚠️ Optional for internal helpers def _internal_helper(data): """Internal - types optional.""" pass ``` ### 2. Use `Any` Strategically ```python # ✅ Good use of Any def process_principle_config(config: dict[str, Any]) -> ConstitutionalPrinciple: """Config structure varies - Any is appropriate.""" pass # ❌ Overuse of Any def add(a: Any, b: Any) -> Any: """Too vague - use specific types.""" return a + b ``` ### 3. Document Complex Types ```python from typing import TypeAlias # Define alias for complex type TrainingBatch: TypeAlias = dict[str, Tensor] EvalResults: TypeAlias = dict[str, float | bool | list[str]] def train_step(batch: TrainingBatch) -> EvalResults: """Type alias makes signature readable.""" pass ``` ### 4. Handle Protocol/Abstract Types ```python from typing import Protocol class Evaluator(Protocol): """Protocol for evaluation functions.""" def evaluate(self, text: str) -> dict[str, Any]: ... def run_evaluation(text: str, evaluator: Evaluator) -> dict[str, Any]: """Accept any object matching Evaluator protocol.""" return evaluator.evaluate(text) ``` ## Mypy Configuration (from pyproject.toml) ```toml [tool.mypy] python_version = "3.10" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = false # Not required for ML code ignore_missing_imports = true # Many ML libraries lack stubs ``` **Key settings:** - `disallow_untyped_defs = false` - Types helpful but not required - `ignore_missing_imports = true` - PyTorch, transformers lack complete type stubs ## When to Use `# type: ignore` ### Acceptable Use Cases ```python # ✅ Known mypy limitation with NumPy scores = np.array(data) # type: ignore[misc] # ✅ PyTorch dynamic typing loss = criterion(outputs, targets) # type: ignore[arg-type] # ✅ Third-party library without stubs from some_ml_lib import advanced_feature # type: ignore ``` ### Avoid Overuse ```python # ❌ Don't ignore fixable errors def add(a: int, b: int) -> str: return a + b # type: ignore # Fix the return type instead! # ✅ Fix the actual issue def add(a: int, b: int) -> int: return a + b ``` ## Type Checking Commands ```bash # Run mypy on all source code mypy constitutional_ai/ --ignore-missing-imports # Run on specific module mypy constitutional_ai/framework.py # Show error codes (useful for targeted ignores) mypy constitutional_ai/ --show-error-codes # Strict mode (educational, will fail) mypy constitutional_ai/ --strict ``` ## Summary: Type Hint Philosophy for ML Code 1. **Use type hints where they add clarity** (public APIs, config) 2. **Skip type hints where they fight the framework** (PyTorch internals) 3. **Accept mypy errors for known ML ecosystem limitations** (documented in MYPY_ANALYSIS_REPORT.md) 4. **Use `Any` strategically, not lazily** (complex structures = OK, simple functions = not OK) 5. **CI should pass with 41 known errors** (continue-on-error: true) 6. **Prioritize correctness over type perfection** (tests are the real validation) ## Quick Reference | Type | When to Use | Example | |------|-------------|---------| | `str`, `int`, `float` | Primitives | `def format(text: str) -> str` | | `list[T]`, `dict[K, V]` | Collections | `def process(data: list[str]) -> dict[str, int]` | | `Tensor` | PyTorch tensors | `def forward(x: Tensor) -> Tensor` | | `Module` | PyTorch models | `def train(model: Module) -> None` | | `PreTrainedModel` | HF models | `def generate(model: PreTrainedModel) -> str` | | `Any` | Unknown/dynamic | `def process(config: dict[str, Any]) -> Any` | | `T \| None` | Optional | `def load(path: str \| None) -> Tensor \| None` | | `tuple[A, B]` | Multiple returns | `def load_model() -> tuple[Module, Tokenizer]` | **Remember**: The goal is helpful type hints, not perfect type coverage. This is research code. --- > Converted and distributed by [TomeVault](https://tomevault.io/claim/apassuello) — claim your Tome and manage your conversions.