pytorch - 💡(How to fix) Fix `torch.compile` produces different output for model with `amax`/`amin` + partial reuse reduction pattern [2 comments, 2 participants]

import os
os.environ["TRITON_BACKENDS_IN_TREE"] = "1"

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiReductionModel(nn.Module):
    def __init__(self, in_channels=3, hidden_dim=64):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(hidden_dim)
        self.relu = nn.ReLU()
        self.conv3 = nn.Conv2d(hidden_dim, hidden_dim // 2, kernel_size=1)
        self.bn2 = nn.BatchNorm2d(hidden_dim // 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        features = self.relu(x)

        # Multiple reductions on the SAME tensor
        dim_amax = torch.amax(features, dim=[2, 3], keepdim=True)
        global_amax = torch.amax(features)
        dim_amin = torch.amin(features, dim=[1], keepdim=True)
        global_amin = torch.min(features)

        # Combine
        amax_normalized = features / (dim_amax + 1e-8)
        amax_scaled = amax_normalized * global_amax
        amin_centered = features - dim_amin
        amin_scaled = amin_centered / (torch.abs(global_amin) + 1e-8)
        combined = amax_scaled + amin_scaled

        output = self.conv3(combined)
        output = self.bn2(output)
        output = self.relu(output)
        return output


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

with torch.no_grad():
    eager_out = model(x)

torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    comp_out = compiled(x)

diff = (eager_out.float() - comp_out.float()).abs().max().item()
print(f"Max diff: {diff}")
print(f"Match: {torch.allclose(eager_out, comp_out, atol=1e-5, rtol=1e-4)}")
# Expected: max_diff ≈ 896.0 — very large mismatch

---

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