pytorch - 💡(How to fix) Fix `torch.compile` produces numerically different results for batch pointwise fusion pattern (`select → unsqueeze → expand → pointwise`) compared to eager mode [1 comments, 1 participants]

import torch
import torch.nn as nn

class SelectParentPointwiseModel(nn.Module):
    def __init__(self, channels=64):
        super().__init__()
        self.conv = nn.Conv2d(3, channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(channels)
        self.scale = nn.Parameter(torch.randn(channels, 1, 1))
        self.shift = nn.Parameter(torch.randn(channels, 1, 1))

    def forward(self, x):
        base = self.bn(self.conv(x))  # [B, C, H, W]

        # Multiple aten.select on channel dim → unsqueeze → expand
        s0 = base.select(1, 0).unsqueeze(1).expand_as(base)
        s1 = base.select(1, 1).unsqueeze(1).expand_as(base)
        s2 = base.select(1, 2).unsqueeze(1).expand_as(base)
        s3 = base.select(1, 3).unsqueeze(1).expand_as(base)

        # Batch pointwise ops: add, mul, div, torch.add with alpha
        p1 = s0 + base                              # add
        p2 = s1 * base                              # mul
        p3 = s2 / (base.abs() + 1e-6)               # div
        p4 = torch.add(s3, base, alpha=0.5)          # add with alpha
        p5 = s0 * self.scale + self.shift            # scale + shift
        p6 = s1 + s2 * base                          # mixed

        combined = p1 + p2 + p3 + p4 + p5 + p6
        return combined

device = "cuda"
torch.manual_seed(42)
model = SelectParentPointwiseModel().to(device).eval()
x = torch.randn(4, 3, 32, 32, device=device)

# Eager: deterministic
with torch.no_grad():
    ref1 = model(x)
    ref2 = model(x)
print(f"Eager deterministic: {(ref1 - ref2).abs().max().item():.6e}")

# Compiled: different result
torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    out = compiled(x)

diff = (ref1.float() - out.float()).abs()
print(f"max_diff={diff.max().item():.6e}")
print(f"Output range: [{ref1.min().item():.2f}, {ref1.max().item():.2f}]")
print(f"Match: {torch.allclose(ref1, out, atol=1e-5, rtol=1e-4)}")

---

Eager deterministic: 0.000000e+00
max_diff > 0 (systematic mismatch from post-grad batch pointwise fusion)

---

PyTorch version: 2.12.0.dev20260327+cu126
Python: 3.10.12
OS: Ubuntu 22.04.5 LTS (WSL2)
GPU: NVIDIA GeForce RTX 3060 Laptop GPU
CUDA: 12.6