import svetoviz_webgpu as sv

# 1. Define Layer Parameters
num_features = 128
bn = nn.BatchNorm1d(num_features=num_features)

# 2. Prepare Input Tensor (Batch=16, Features=128)
input_tensor = torch.randn(16, num_features) * 5.0 + 2.0

def terminal_callback(buffer, message, images, files):
    bn.train()  # Set to training mode to update running stats
    output = bn(input_tensor)

    buffer.send_system_message(f"Input Shape: {list(input_tensor.shape)}")
    buffer.send_system_message(f"Running Mean (First 5): {bn.running_mean[:5].tolist()}")
    buffer.send_system_message("Status: Normalized across the batch dimension (N).")

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