Automate PyTorch Unit Tests with AI in 20 Minutes

# Install required packages
pip install torch pytest anthropic --break-system-packages

# Verify installation
python -c "import torch; import anthropic; print('Ready')"

# test_generator.py
import anthropic
import os
from pathlib import Path

def generate_tests(model_code: str, model_name: str) -> str:
    """
    Generate pytest tests for PyTorch model using Claude API.
    
    Args:
        model_code: Source code of the PyTorch model
        model_name: Name of the model class
    
    Returns:
        Generated test code as string
    """
    client = anthropic.Anthropic(
        api_key=os.environ.get("ANTHROPIC_API_KEY")
    )
    
    # This prompt is optimized for PyTorch-specific test generation
    prompt = f"""Generate comprehensive pytest unit tests for this PyTorch model.

MODEL CODE:
{model_code}

REQUIREMENTS:
1. Test tensor shape consistency across forward pass
2. Test gradient flow (no NaN, no vanishing)
3. Test with different batch sizes (1, 16, 32)
4. Test device compatibility (CPU/CUDA if available)
5. Test edge cases (zero input, negative values, extreme values)
6. Use fixtures for model instantiation
7. Include docstrings explaining what each test validates

MODEL CLASS: {model_name}

Return ONLY the pytest code, no explanations."""

    message = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=4000,
        messages=[{"role": "user", "content": prompt}]
    )
    
    return message.content[0].text

def save_tests(test_code: str, output_path: str):
    """Save generated tests to file"""
    Path(output_path).write_text(test_code)
    print(f"✓ Tests saved to {output_path}")

if __name__ == "__main__":
    # Example usage
    sample_model = """
import torch
import torch.nn as nn

class SimpleClassifier(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=128, num_classes=10):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, num_classes)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x
"""
    
    tests = generate_tests(sample_model, "SimpleClassifier")
    save_tests(tests, "test_classifier.py")

export ANTHROPIC_API_KEY="your-api-key-here"

# your_model.py
import torch
import torch.nn as nn

class CustomCNN(nn.Module):
    """Example: Image classification CNN"""
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(64 * 8 * 8, 10)
    
    def forward(self, x):
        # x shape: (batch, 3, 32, 32)
        x = self.pool(torch.relu(self.conv1(x)))  # (batch, 32, 16, 16)
        x = self.pool(torch.relu(self.conv2(x)))  # (batch, 64, 8, 8)
        x = x.view(x.size(0), -1)  # Flatten
        x = self.fc(x)
        return x

python -c "
from test_generator import generate_tests, save_tests
from pathlib import Path

model_code = Path('your_model.py').read_text()
tests = generate_tests(model_code, 'CustomCNN')
save_tests(tests, 'test_custom_cnn.py')
"

# test_custom_cnn.py (AI-generated example)
import pytest
import torch
from your_model import CustomCNN

@pytest.fixture
def model():
    """Fixture providing fresh model instance for each test"""
    return CustomCNN()

def test_forward_pass_shape(model):
    """Verify output shape matches expected num_classes"""
    batch_size = 16
    x = torch.randn(batch_size, 3, 32, 32)
    output = model(x)
    
    assert output.shape == (batch_size, 10), \
        f"Expected (16, 10), got {output.shape}"

def test_gradient_flow(model):
    """Ensure gradients propagate without NaN or explosion"""
    x = torch.randn(4, 3, 32, 32, requires_grad=True)
    output = model(x)
    loss = output.sum()
    loss.backward()
    
    # Check no NaN gradients
    for name, param in model.named_parameters():
        assert not torch.isnan(param.grad).any(), \
            f"NaN gradient in {name}"
        assert not torch.isinf(param.grad).any(), \
            f"Inf gradient in {name}"

def test_batch_size_flexibility(model):
    """Model should handle different batch sizes"""
    for batch_size in [1, 16, 32]:
        x = torch.randn(batch_size, 3, 32, 32)
        output = model(x)
        assert output.shape[0] == batch_size

def test_zero_input_handling(model):
    """Model should process zero tensors without errors"""
    x = torch.zeros(8, 3, 32, 32)
    output = model(x)
    assert not torch.isnan(output).any()

# If AI assumes wrong input shape, update the test:
# Change this:
x = torch.randn(batch_size, 3, 32, 32)
# To your actual input shape:
x = torch.randn(batch_size, 1, 28, 28)  # MNIST example

# Run all generated tests
pytest test_custom_cnn.py -v

# Run with coverage report
pytest test_custom_cnn.py --cov=your_model --cov-report=term-missing

# 1. Generate tests for a new model
python test_generator.py

# 2. Run generated tests
pytest test_classifier.py -v

# 3. Check coverage
pytest --cov=. --cov-report=html
open htmlcov/index.html  # View coverage report

name: Auto-Generated Tests

on: [push, pull_request]

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      
      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      
      - name: Install dependencies
        run: |
          pip install torch pytest anthropic --break-system-packages
      
      - name: Generate tests
        env:
          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
        run: python test_generator.py
      
      - name: Run tests
        run: pytest -v --cov=. --cov-report=xml
      
      - name: Upload coverage
        uses: codecov/codecov-action@v4

# test_templates.py
TEMPLATE_PROMPT = """Generate pytest tests for this PyTorch model.

CRITICAL REQUIREMENTS:
- Test on our standard input shape: (batch, 512, 768)
- Verify output is probability distribution (sum to 1)
- Test with our custom loss function: FocalLoss
- Include tests for mixed precision (fp16)

MODEL CODE:
{model_code}

Return pytest code only."""

# batch_generate.py
from pathlib import Path
from test_generator import generate_tests, save_tests

models_dir = Path("models/")
for model_file in models_dir.glob("*.py"):
    code = model_file.read_text()
    tests = generate_tests(code, model_file.stem)
    save_tests(tests, f"tests/test_{model_file.stem}.py")

# Only regenerate tests when model changes
import hashlib

def should_regenerate(model_path, test_path):
    """Check if model changed since last test generation"""
    model_hash = hashlib.sha256(Path(model_path).read_bytes()).hexdigest()
    
    if not Path(test_path).exists():
        return True
    
    # Store hash in test file comment
    test_content = Path(test_path).read_text()
    if f"# model_hash: {model_hash}" in test_content:
        return False
    return True

# Debug: Check if model imports correctly
python -c "from your_model import CustomCNN; print('OK')"

class CustomCNN(nn.Module):
    """
    Image classifier for 32x32 RGB images.
    
    Input shape: (batch_size, 3, 32, 32)
    Output shape: (batch_size, 10)
    """

# Add retry logic
from anthropic import RateLimitError
import time

def generate_with_retry(model_code, model_name, max_retries=3):
    for attempt in range(max_retries):
        try:
            return generate_tests(model_code, model_name)
        except RateLimitError:
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # Exponential backoff
            else:
                raise