import torch

d_model = 512
vocab_size = 10000
batch_size = 4

x = torch.randn(batch_size, d_model, device="cuda")
weight = torch.randn(vocab_size, d_model, device="cuda")
bias = torch.zeros(d_model, device="cuda")  # shape [512], not [10000]

# Eager: succeeds because beta=0.0 zeros out the bias term
try:
    out = torch.addmm(bias, x, weight.t(), beta=0.0, alpha=0.1)
    print(f"eager: OK shape={out.shape}")
except RuntimeError as e:
    print(f"eager: ERROR — {e}")

# Compiled: raises shape mismatch error
torch._dynamo.reset()

@torch.compile(fullgraph=True)
def compiled_addmm(bias, x, weight):
    return torch.addmm(bias, x, weight.t(), beta=0.0, alpha=0.1)

try:
    out = compiled_addmm(bias, x, weight)
    print(f"compile: OK shape={out.shape}")
except RuntimeError as e:
    print(f"compile: ERROR — {e}")

---

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

class TransformerStyleModel(nn.Module):
    def __init__(self, d_model=512, nhead=8, num_layers=3, vocab_size=10000, max_seq_len=512):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = nn.Parameter(torch.randn(1, max_seq_len, d_model) * 0.02)
        self.ffn_weight1 = nn.Parameter(torch.randn(d_model * 4, d_model) * 0.02)
        self.ffn_bias1 = nn.Parameter(torch.zeros(d_model * 4))
        self.ffn_weight2 = nn.Parameter(torch.randn(d_model, d_model * 4) * 0.02)
        self.ffn_bias2 = nn.Parameter(torch.zeros(d_model))
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.output_proj = nn.Linear(d_model, vocab_size)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        batch_size, seq_len = x.shape
        x = self.embedding(x) * self.d_model ** 0.5
        x = x + self.pos_encoding[:, :seq_len, :]
        x = self.dropout(x)
        x_flat = x.view(-1, self.d_model)
        ffn_intermediate = torch.addmm(self.ffn_bias1, x_flat, self.ffn_weight1.t(), beta=1.0, alpha=1.0)
        ffn_intermediate = F.gelu(ffn_intermediate)
        ffn_intermediate = ffn_intermediate * 0.5
        ffn_intermediate = self.dropout(ffn_intermediate)
        ffn_output = torch.addmm(self.ffn_bias2, ffn_intermediate, self.ffn_weight2.t(), beta=1.0, alpha=1.0)
        ffn_output = torch.tanh(ffn_output)
        ffn_output = ffn_output + x_flat
        ffn_output = self.ln1(ffn_output)
        ffn_output = ffn_output.view(batch_size, seq_len, self.d_model)
        x_flat2 = ffn_output.view(-1, self.d_model)
        # BUG: bias shape [512] mismatches output [10000], but beta=0.0 makes it work in eager
        logits = torch.addmm(torch.zeros(self.d_model, device=x.device), x_flat2,
                              self.output_proj.weight.t(), beta=0.0, alpha=0.1)
        logits = logits + self.output_proj.bias
        logits = torch.softmax(logits, dim=-1)
        logits = logits * 0.9 + 0.05
        logits = logits.view(batch_size, seq_len, -1)
        return logits

model = TransformerStyleModel().cuda().eval()
x = torch.randint(0, 10000, (2, 32), dtype=torch.long, device="cuda")

# Eager: succeeds
try:
    with torch.no_grad():
        out = model(x)
    print(f"eager: OK shape={out.shape}")
except RuntimeError as e:
    print(f"eager: ERROR — {e}")

# Compiled: fails
torch._dynamo.reset()
compiled_model = torch.compile(model, fullgraph=True)
try:
    with torch.no_grad():
        out = compiled_model(x)
    print(f"compile: OK shape={out.shape}")
except RuntimeError as e:
    print(f"compile: ERROR — {e}")

---

(no error — returns tensor of shape [batch*seq, vocab_size])

---

RuntimeError: The size of tensor a (512) must match the size of tensor b (10000)
at non-singleton dimension 1

---

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