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.bn1 = nn.BatchNorm2d(hidden_dim)
        self.conv2 = nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1)
        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 (triggers reuse_partial)
        spatial_amax = torch.amax(features, dim=[2, 3], keepdim=True)   # per-channel spatial max
        global_amax = torch.amax(features)                               # global max scalar
        channel_amin = torch.amin(features, dim=[1], keepdim=True)       # per-spatial-loc channel min
        global_min_val = torch.min(features)                             # global min scalar

        # Normalize and center using reduction results
        amax_normalized = features / (spatial_amax + 1e-8)
        amax_scaled = amax_normalized * global_amax
        amin_centered = features - channel_amin
        amin_scaled = amin_centered / (torch.abs(global_min_val) + 1e-8)

        # Combine normalized and centered features
        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)

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

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

diff = (ref.float() - comp.float()).abs()
print(f"Max diff: {diff.max().item():.6f}")
print(f"Mean diff: {diff.mean().item():.6e}")
print(f"Match (atol=1e-5): {torch.allclose(ref, comp, atol=1e-5, rtol=1e-4)}")
# Expected: max_diff in the hundreds — very large mismatch

---

Eager deterministic: 0.000000e+00
Max diff: ~896.0+ (very large systematic mismatch)

---

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