import torch
from torch import nn


class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(8, 16),
            nn.ReLU(),
            nn.Linear(16, 3),
        )

    def forward(self, inputs):
        return self.layers(inputs)


torch.manual_seed(7)
model = TinyClassifier().eval()
sample = torch.randn(4, 8)
compiled_model = torch.compile(model, backend="inductor")

with torch.inference_mode():
    eager_output = model(sample)
    compiled_output = compiled_model(sample)
    second_compiled_output = compiled_model(sample)

max_abs_diff = (eager_output - compiled_output).abs().max().item()

print(f"torch={torch.__version__}")
print(f"input_shape={tuple(sample.shape)}")
print(f"eager_shape={tuple(eager_output.shape)}")
print(f"compiled_shape={tuple(compiled_output.shape)}")
print(f"second_run_shape={tuple(second_compiled_output.shape)}")
print(f"max_abs_diff={max_abs_diff:.8f}")
print(f"outputs_match={torch.allclose(eager_output, compiled_output, atol=1e-6)}")