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

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

class DiffusionStyleModel(nn.Module):
    def __init__(self, in_channels=3, hidden_channels=64, num_timesteps=1000):
        super().__init__()
        self.num_timesteps = num_timesteps

        # Pointwise (1x1) conv layers
        self.conv_pw1 = nn.Conv2d(in_channels, hidden_channels, kernel_size=1)
        self.bn1 = nn.BatchNorm2d(hidden_channels)
        self.conv_pw2 = nn.Conv2d(hidden_channels, hidden_channels, kernel_size=1)
        self.bn2 = nn.BatchNorm2d(hidden_channels)
        self.conv_pw3 = nn.Conv2d(hidden_channels, in_channels, kernel_size=1)

        # Timestep embedding
        self.time_embed = nn.Sequential(
            nn.Linear(1, hidden_channels),
            nn.GELU(),
            nn.Linear(hidden_channels, hidden_channels),
        )

    def _gelu_erf(self, x):
        """GELU via erf approximation."""
        return 0.5 * x * (1.0 + torch.erf(x * 0.7071067811865476))

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

        # randint #1: sample random timesteps per batch element
        t = torch.randint(0, self.num_timesteps, (B,), device=x.device, dtype=torch.long)
        t_float = t.float() / self.num_timesteps  # normalize to [0, 1)
        t_embed = self.time_embed(t_float.unsqueeze(-1))  # [B, hidden]
        t_embed = t_embed.unsqueeze(-1).unsqueeze(-1)      # [B, hidden, 1, 1]

        # Feature extraction with pointwise convs
        h = self.conv_pw1(x)
        h = self.bn1(h)
        h = self._gelu_erf(h)

        # Add timestep embedding
        h = h + t_embed

        h = self.conv_pw2(h)
        h = self.bn2(h)
        h = self._gelu_erf(h)

        # randint #2: noise mask (per-sample probability mask)
        noise_prob = torch.randint(0, 100, (B, 1, 1, 1), device=x.device, dtype=torch.int32)
        noise_mask = (noise_prob < 50).float()  # ~50% chance per sample
        h = h * noise_mask

        # randint #3: stochastic depth (binary per-sample)
        stoch_depth = torch.randint(0, 2, (B, 1, 1, 1), device=x.device, dtype=torch.int32).float()
        h = h * stoch_depth

        # Output projection
        out = self.conv_pw3(h)

        # Residual connection
        out = x + out

        # randint #4: additive integer noise with negative range
        int_noise = torch.randint(-5, 6, (B, C, H, W), device=x.device, dtype=torch.int32).float()
        out = out + int_noise * 0.01

        return out


device = "cuda"
torch.manual_seed(42)
model = DiffusionStyleModel(
    in_channels=3, hidden_channels=64, num_timesteps=1000
).to(device).eval()
x = torch.randn(4, 3, 32, 32, device=device)

# Eager: runs successfully
with torch.no_grad():
    eager_out = model(x)
print(f"Eager output shape: {eager_out.shape}")
print(f"Eager output range: [{eager_out.min().item():.4f}, {eager_out.max().item():.4f}]")
print("Eager: OK")

# Compiled: raises error
torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
try:
    with torch.no_grad():
        comp_out = compiled(x)
    # If it doesn't crash, check for status difference
    diff = (eager_out.float() - comp_out.float()).abs()
    print(f"Compiled output shape: {comp_out.shape}")
    print(f"Compiled max_diff: {diff.max().item():.6e}")
    print("NOTE: randint outputs are non-deterministic, so diff is expected.")
    print("Check for shape/dtype mismatches or runtime errors instead.")
except Exception as e:
    print(f"torch.compile FAILED: {type(e).__name__}: {e}")

---

Traceback (most recent call last):
  File "reproducer.py", line 82, in <module>
    comp_out = compiled(x)
  ...
  File ".../torch/_inductor/fx_passes/joint_graph.py", line ..., in replace_randint
    ...
  File ".../torch/_inductor/lowering.py", line ..., in ...
RuntimeError: failed to lower randint with low=-5 in fused kernel:
  Inductor's replace_randint optimization does not support negative low values
  in randint(-5, 6, ...) when fused with other randint calls in the same graph

---

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