example-1

I'm an Engineering Director at Google working on GKE. I'm a maintainer of Kubernetes OSS and chair of SIG API-Machinery. I'm expanding my team's OSS influence into the PyTorch ecosystem, starting with TPU support contributions to https://github.com/pytorch/examples.

## What we've already done

In a previous session we researched the pytorch/examples repo and identified 7 candidates for adding TPU/XLA support. We're starting with the MNIST example (`mnist/main.py`) as the first PR.

We've already drafted the code changes. Here's what needs to happen:

## Task

1. Fork and clone `pytorch/examples`
2. Create a branch called `add-xla-tpu-support-mnist`
3. Apply the following changes to `mnist/main.py` and `mnist/README.md`
4. Run linting (`flake8`) and fix any issues
5. Prepare the commit

## Changes to `mnist/main.py`

Replace the entire file with this:

```python
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR


# Optional: XLA/TPU support
_XLA_AVAILABLE = False
try:
    import torch_xla
    import torch_xla.core.xla_model as xm
    _XLA_AVAILABLE = True
except ImportError:
    pass


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        # On XLA devices, use xm.optimizer_step to sync gradients properly.
        # On other devices, use the standard optimizer.step().
        if args.xla:
            xm.optimizer_step(optimizer)
        else:
            optimizer.step()
        if batch_idx % args.log_interval == 0:
            # On XLA, .item() triggers a device-to-host sync. We accept
            # this cost at logging boundaries for readability.
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=14, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-accel', action='store_true',
                        help='disables accelerator')
    parser.add_argument('--xla', action='store_true', default=False,
                        help='enables XLA device (e.g. TPU). Requires torch_xla.')
    parser.add_argument('--dry-run', action='store_true',
                        help='quickly check a single pass')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true',
                        help='For Saving the current Model')
    args = parser.parse_args()

    # Device selection: --xla takes precedence, then torch.accelerator, then CPU.
    if args.xla:
        if not _XLA_AVAILABLE:
            raise RuntimeError(
                "--xla flag requires torch_xla to be installed. "
                "Install with: pip install torch_xla[tpu]"
            )
        device = xm.xla_device()
        print(f"Using XLA device: {device}")
    else:
        use_accel = not args.no_accel and torch.accelerator.is_available()
        if use_accel:
            device = torch.accelerator.current_accelerator()
        else:
            device = torch.device("cpu")

    torch.manual_seed(args.seed)

    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}

    if args.xla:
        # On XLA/TPU: pin_memory is not needed (no host-device DMA path),
        # and persistent_workers helps avoid re-forking.
        xla_kwargs = {'num_workers': 4,
                      'persistent_workers': True,
                      'shuffle': True,
                      'drop_last': True}
        train_kwargs.update(xla_kwargs)
        # For test, don't drop_last so we evaluate every sample.
        test_kwargs.update({'num_workers': 4, 'persistent_workers': True})
    elif not args.no_accel and torch.accelerator.is_available():
        accel_kwargs = {'num_workers': 1,
                        'persistent_workers': True,
                        'pin_memory': True,
                        'shuffle': True}
        train_kwargs.update(accel_kwargs)
        test_kwargs.update(accel_kwargs)

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()
        # On XLA, explicitly sync at epoch boundaries to ensure all
        # device computations have completed.
        if args.xla:
            xm.mark_step()

    if args.save_model:
        # On XLA devices, move the model to CPU before saving for portability.
        if args.xla:
            model_to_save = model.cpu()
            torch.save(model_to_save.state_dict(), "mnist_cnn.pt")
            model.to(device)  # move back if needed
        else:
            torch.save(model.state_dict(), "mnist_cnn.pt")


if __name__ == '__main__':
    main()
```

## Changes to `mnist/README.md`

Replace the entire file with this:

```markdown
# MNIST Example

Trains a ConvNet on the MNIST dataset using PyTorch.

## Usage

### Standard (CUDA / MPS / XPU)

\```bash
pip install -r requirements.txt
python main.py
# or to run on CPU only:
python main.py --no-accel
\```

### TPU (via PyTorch/XLA)

To train on a Google Cloud TPU VM:

\```bash
# Install PyTorch/XLA (adjust version as needed)
pip install torch torchvision
pip install 'torch_xla[tpu]'

# Run with the --xla flag
python main.py --xla
\```

**Notes on TPU training:**
- The `--xla` flag selects the XLA device (TPU). It takes precedence over `torch.accelerator`.
- `drop_last=True` is used for the training DataLoader to ensure fixed batch shapes, which avoids recompilation on TPU.
- Model checkpoints saved with `--save-model` are moved to CPU before saving for cross-device portability.
- For multi-device TPU training (e.g. all 4 chips on a v4-8), see the [PyTorch/XLA multiprocessing guide](https://docs.pytorch.org/xla/master/learn/pytorch-on-xla-devices.html).

### Options

\```
usage: main.py [-h] [--batch-size N] [--test-batch-size N] [--epochs N]
               [--lr LR] [--gamma M] [--no-accel] [--xla] [--dry-run]
               [--seed S] [--log-interval N] [--save-model]

PyTorch MNIST Example

options:
  -h, --help           show this help message and exit
  --batch-size N       input batch size for training (default: 64)
  --test-batch-size N  input batch size for testing (default: 1000)
  --epochs N           number of epochs to train (default: 14)
  --lr LR              learning rate (default: 1.0)
  --gamma M            Learning rate step gamma (default: 0.7)
  --no-accel           disables accelerator
  --xla                enables XLA device (e.g. TPU). Requires torch_xla.
  --dry-run            quickly check a single pass
  --seed S             random seed (default: 1)
  --log-interval N     how many batches to wait before logging training status
  --save-model         For Saving the current Model
\```
```

## PR description to use

Title: `Add XLA/TPU support to MNIST example`

Body:

```
## Summary

Adds Google Cloud TPU support to the MNIST example via PyTorch/XLA, while
preserving full backward compatibility with the existing CUDA/MPS/XPU paths.
No model architecture changes — the same ConvNet trains identically on TPU.

## Changes

### `mnist/main.py`
- **Optional `torch_xla` import**: Guarded `try/except` so the example still
  works without `torch_xla` installed.
- **`--xla` flag**: New CLI argument to opt in to XLA/TPU device. Takes
  precedence over `torch.accelerator`.
- **`xm.optimizer_step()`**: Replaces `optimizer.step()` on XLA to properly
  synchronize gradients across the XLA computation graph.
- **DataLoader config**: TPU-specific settings — `drop_last=True` for training
  (avoids recompilation from ragged last batch), `num_workers=4`, no
  `pin_memory` (not applicable to XLA).
- **`xm.mark_step()`** at epoch boundaries for explicit sync.
- **CPU checkpoint save**: Moves model to CPU before `torch.save()` for
  cross-device portability.

### `mnist/README.md`
- Added "TPU (via PyTorch/XLA)" section with setup and usage instructions.

## Testing

- [x] CPU (no flags) — behavior unchanged
- [x] `--no-accel` — behavior unchanged
- [ ] CUDA — needs CI verification
- [ ] TPU v4-8 with `--xla` — trains to ~99% accuracy in 14 epochs

## Dependencies

`torch_xla` is optional. The example works without it unless `--xla` is passed.
Minimum version: `torch_xla >= 2.6.0`.

This is the first in a planned series adding TPU support to other examples
(VAE, word_language_model, super_resolution, dcgan, imagenet).
```

## Instructions

1. Fork `pytorch/examples` to my GitHub account
2. Clone it locally
3. Create branch `add-xla-tpu-support-mnist`
4. Write the files above to `mnist/main.py` and `mnist/README.md`
5. Run `flake8 mnist/main.py` and fix any lint issues
6. Also verify the file runs on CPU: `python mnist/main.py --dry-run --no-accel`
7. Commit with message: `Add XLA/TPU support to MNIST example`
8. Show me the final diff before pushing