import svetoviz_webgpu as sv

# 1. Define Layer Parameters
features = 128
ln = nn.LayerNorm(normalized_shape=features)

# 2. Prepare Input Tensor (Batch=2, Features=128)
# High variance and shifted mean to demonstrate normalization
input_tensor = torch.randn(2, features) * 10 + 5.0

def terminal_callback(buffer, message, images, files):
    # Forward pass
    output = ln(input_tensor)

    # Calculate stats for the first sample
    mean_out = output[0].mean().item()
    var_out = output[0].var(unbiased=False).item()

    buffer.send_system_message(f"Normalization Dim: {features}")
    buffer.send_system_message(f"Sample 0 Mean (Post-Norm): {mean_out:.4f}")
    buffer.send_system_message(f"Sample 0 Variance (Post-Norm): {var_out:.4f}")
    buffer.send_system_message("Status: Normalized across the feature axis per-sample.")

# 3. Start the interactive session
sv.pytorch_web(module=ln, terminal_callback=terminal_callback)